import Wayland
import Waylean.Registry
import Std.Data.HashMap

open Wayland Protocols Waylean

def Std.HashMap.mapM {α β γ : Type} [BEq α] [Hashable α] {m : Type → Type} [Monad m] (map : Std.HashMap α β) (update : α → β → m γ) : m (Std.HashMap α γ) := do
  let a := map.toArray
  let a' ← a.mapM fun (x, y) ↦ do return (x, ← update x y)
  return Std.HashMap.ofArray a'

def Std.HashMap.modifyM {α β : Type} [BEq α] [Hashable α] {m : Type → Type} [Monad m] (map : Std.HashMap α β) (key : α) (update : β → m β) : m (Std.HashMap α β) := do
  if h : map.contains key then
    let new ← update (map.get key h)
    return map.insert key new
  else return map

inductive Action where
  | none
  | spawnFoot
  | close
  | focusNext
  | move
  | resize
  | exit

inductive SeatOp where
  | none
  | move (window : RiverWindowV1) (startX startY : Int32)
  | resize (window : RiverWindowV1) (startX startY startWidth startHeight : Int32) (edges : RiverWindowV1.Edges)

structure XkbBinding where
  proxy : RiverXkbBindingV1
  action : Action

structure PointerBinding where
  proxy : RiverPointerBindingV1
  action : Action

structure Window where
  proxy : RiverWindowV1
  node : RiverNodeV1
  new : Bool := true
  closed : Bool := false
  x : Int32 := 0
  y : Int32 := 0
  width : Int32 := 0
  height : Int32 := 0
  pointerMoveRequested : Option RiverSeatV1 := .none
  pointerResizeRequested : Option RiverSeatV1 := .none
  pointerResizeRequestedEdges : RiverWindowV1.Edges := {}

structure Output where
  proxy : RiverOutputV1
  removed : Bool := false

structure Seat where
  proxy : RiverSeatV1
  new : Bool := true
  removed : Bool := false
  focused : Option RiverWindowV1 := .none
  hovered : Option RiverWindowV1 := .none
  interacted : Option RiverWindowV1 := .none
  xkbBindings : Std.HashMap RiverXkbBindingV1 XkbBinding := {}
  pointerBindings : Std.HashMap RiverPointerBindingV1 PointerBinding := {}
  pendingAction : Action := .none
  op : SeatOp := .none
  opDeltaX : Int32 := 0
  opDeltaY : Int32 := 0
  opRelease : Bool := false

structure State where
  registry : Registry.State := .new
  windows : Std.HashMap RiverWindowV1 Window := {}
  windowOrder : Array RiverWindowV1 := #[]
  outputs : Std.HashMap RiverOutputV1 Output := {}
  seats : Std.HashMap RiverSeatV1 Seat := {}

def zoomRegistry {α : Type} : DisplayM Registry.State α → DisplayM State α := DisplayM.zoom (mk_lens registry)

def handleSeat (ev : RiverSeatV1.Event) : DisplayM (Option Seat) Unit := do
  for s in ← get do
    match ev with
      | .removed               => set (Option.some { s with removed := true })
      | .pointerEnter win      => set (Option.some { s with hovered := win })
      | .pointerLeave          => set (Option.some { s with hovered := .none })
      | .windowInteraction win => set (Option.some { s with interacted := win })
      | .opDelta x y           => set (Option.some { s with opDeltaX := x, opDeltaY := y })
      | .opRelease             => set (Option.some { s with opRelease := true })
      | _                      => pure ()

def handleXkbBinding (proxy : RiverXkbBindingV1) (ev : RiverXkbBindingV1.Event) : DisplayM (Option Seat) Unit := do
  for s in ← get do
    for binding in s.xkbBindings.get? proxy do
      match ev with
        | .pressed => set (Option.some { s with pendingAction := binding.action })
        | _        => pure ()

def handlePointerBinding (proxy : RiverPointerBindingV1) (ev : RiverPointerBindingV1.Event) : DisplayM (Option Seat) Unit := do
  for s in ← get do
    for binding in s.pointerBindings.get? proxy do
      match ev with
        | .pressed => set (Option.some { s with pendingAction := binding.action })
        | _        => pure ()

def handleOutput (ev : RiverOutputV1.Event) : DisplayM (Option Output) Unit := do
  for o in ← get do
    match ev with
      | .removed => set (Option.some { o with removed := true })
      | _        => pure ()

def handleWindow (ev : RiverWindowV1.Event) : DisplayM (Option Window) Unit := do
  for w in ← get do
    match ev with
      | .closed                     => set (Option.some { w with closed := true })
      | .dimensions width height    => set (Option.some { w with width, height })
      | .pointerMoveRequested s     => set (Option.some { w with pointerMoveRequested := s })
      | .pointerResizeRequested s e => set (Option.some { w with pointerResizeRequested := s, pointerResizeRequestedEdges := e })
      | _                           => pure ()

structure Context where
  wm : RiverWindowManagerV1
  xkbBindings : RiverXkbBindingsV1

def Window.setPosition (x y : Int32) : DisplayM (Option Window) Unit := do
  for w in ← get do
    w.node.setPosition x y
    set (Option.some { w with x, y})

def Seat.createXkbBinding (ctx : Context) (mods : RiverSeatV1.Modifiers) (keysym : UInt32) (action : Action) : DisplayM (Option Seat) Unit := do
  for s in ← get do
    let proxy ← ctx.xkbBindings.getXkbBinding s.proxy keysym mods
    proxy.enable
    proxy.listen (handleXkbBinding proxy)
    let xkbBindings := s.xkbBindings.insert proxy { proxy, action }
    set (Option.some { s with xkbBindings })

def Seat.createPointerBinding (mods : RiverSeatV1.Modifiers) (button : UInt32) (action : Action) : DisplayM (Option Seat) Unit := do
  for s in ← get do
    let proxy ← s.proxy.getPointerBinding button mods
    proxy.enable
    proxy.listen (handlePointerBinding proxy)
    let pointerBindings := s.pointerBindings.insert proxy { proxy, action }
    set (Option.some { s with pointerBindings })

def Seat.pointerMove {σ : Type} (seat : Seat) (win : Window) : DisplayM σ Seat := do
  seat.proxy.opStartPointer
  return { seat with
    interacted := win.proxy
    op := .move win.proxy win.x win.y
    opDeltaX := 0
    opDeltaY := 0
  }

def Seat.pointerResize {σ : Type} (seat : Seat) (win : Window) (edges : RiverWindowV1.Edges) : DisplayM σ Seat := do
  seat.proxy.opStartPointer
  win.proxy.informResizeStart
  return { seat with
    interacted := win.proxy
    op := .resize win.proxy win.x win.y win.width win.height edges
    opDeltaX := 0
    opDeltaY := 0
  }

def Seat.focusTop (s : Seat) : DisplayM State Seat := do
  let windows := (← get).windowOrder
  if h : windows.size != 0 then
    let proxy := windows[windows.size - 1]'(Nat.sub_one_lt (bne_iff_ne.mp h))
    for win in (← get).windows.get? proxy do
      s.proxy.focusWindow win.proxy
      win.node.placeTop
      return { s with focused := win.proxy }
  s.proxy.clearFocus
  return { s with focused := .none }

def Seat.doAction (s : Seat) (ctx : Context) : DisplayM State Seat := do
  let action := s.pendingAction
  let mut s : Seat := { s with pendingAction := .none }
  match action with
    | .none =>
        pure ()
    | .spawnFoot =>
        _ ← DisplayM.liftIO <| IO.Process.spawn {
          cmd := "foot"
          env := #[("WAYLAND_DEBUG", .none)]
        }
    | .close =>
        for window in s.focused do
          window.close
    | .focusNext =>
        let mut windowOrder := (← get).windowOrder
        if h : 0 < windowOrder.size then
          let win := windowOrder[0]
          windowOrder := windowOrder.eraseIdx 0
          windowOrder := windowOrder.push win
        set { ← get with windowOrder }
        s ← focusTop s
    | .move =>
        for win in s.hovered >>= (← get).windows.get? do
          s ← s.pointerMove win
    | .resize =>
        for win in s.hovered >>= (← get).windows.get? do
          s ← s.pointerResize win { bottom := true, right := true }
    | .exit =>
        ctx.wm.exitSession
  return s

def Seat.opEnd (s : Seat) : DisplayM State Seat := do
  match s.op with
    | .resize proxy _ _ _ _ _ => proxy.informResizeEnd
    | _ => pure ()
  s.proxy.opEnd
  return { s with op := .none }

def Seat.opManage (s : Seat) : DisplayM State Unit := do
  match s.op with
    | .resize proxy _ _ startWidth startHeight edges =>
        let mut width := startWidth
        let mut height := startHeight
        if edges.left   then width  := width  - s.opDeltaX
        if edges.right  then width  := width  + s.opDeltaX
        if edges.top    then height := height - s.opDeltaY
        if edges.bottom then height := height + s.opDeltaY
        width  := if width  > 0 then width  else 1
        height := if height > 0 then height else 1
        proxy.proposeDimensions width height
    | _ => pure ()

def removeOutputs : DisplayM State Unit := do
  let (removedOutputs, aliveOutputs) := (← get).outputs.partition fun _ out ↦ out.removed
  for out in removedOutputs.values do
    out.proxy.destroy
  set { ← get with outputs := aliveOutputs }

def removeWindows : DisplayM State Unit := do
  let (closedWindows, aliveWindows) := (← get).windows.partition fun _ win ↦ win.closed
  let mut seats := (← get).seats
  let mut windowOrder := (← get).windowOrder
  for win in closedWindows.values do
    windowOrder := windowOrder.erase win.proxy
    seats ← seats.mapM fun _ seat ↦ do
      let op  ← match seat.op with
        | .none => pure SeatOp.none
        | op@(.move proxy _ _) =>
            if proxy == win.proxy then
              seat.proxy.opEnd
              pure SeatOp.none
            else pure op
        | op@(.resize proxy _ _ _ _ _) =>
            if proxy == win.proxy then
              seat.proxy.opEnd
              pure SeatOp.none
            else pure op
      pure { seat with
        op
        focused := if seat.focused == win.proxy then .none else seat.focused
        hovered := if seat.hovered == win.proxy then .none else seat.hovered
        interacted := if seat.interacted == win.proxy then .none else seat.interacted
      }
    win.proxy.destroy
  set { ← get with windows := aliveWindows, windowOrder, seats }

def removeSeats : DisplayM State Unit := do
  let (removedSeats, aliveSeats) := (← get).seats.partition fun _ seat ↦ seat.removed
  for seat in removedSeats.values do
    for binding in seat.xkbBindings.values do
      binding.proxy.destroy
    for binding in seat.pointerBindings.values do
      binding.proxy.destroy
    seat.proxy.destroy
  set { ← get with seats := aliveSeats }

def initNewWindows : DisplayM State Unit := do
  for w in (← get).windows.values.filter Window.new do
    DisplayM.zoom (mk_lens windows ≫ Std.HashMap.lens w.proxy) do
      set (Option.some { w with new := false })
      Window.setPosition w.x w.y
      w.proxy.proposeDimensions w.width w.height

def initNewSeats (ctx : Context) : DisplayM State Unit := do
  -- See xkbcommon/xkbcommon-keysyms.h
  let space : UInt32 := 0x20
  let n : UInt32 := 0x6e
  let q : UInt32 := 0x71
  let esc : UInt32 := 0xff1b
  -- See linux/input-event-codes.h
  let btnLeft : UInt32 := 0x110
  let btnRight : UInt32 := 0x111
  let mods : RiverSeatV1.Modifiers := { mod1 := true }

  for s in (← get).seats.values.filter Seat.new do
    DisplayM.zoom (mk_lens seats ≫ Std.HashMap.lens s.proxy) do
      set (Option.some { s with new := false })
      Seat.createXkbBinding ctx mods space .spawnFoot
      Seat.createXkbBinding ctx mods q .close
      Seat.createXkbBinding ctx mods n .focusNext
      Seat.createXkbBinding ctx mods esc .exit
      Seat.createPointerBinding mods btnLeft .move
      Seat.createPointerBinding mods btnRight .resize

def manageWindows : DisplayM State Unit := do
  let windows ← (← get).windows.mapM fun _ win ↦ do
    let mut seats := (← get).seats
    for proxy in win.pointerMoveRequested do
      seats ← seats.modifyM proxy (fun s ↦ s.pointerMove win)
    for proxy in win.pointerResizeRequested do
      seats ← seats.modifyM proxy (fun s ↦ s.pointerResize win win.pointerResizeRequestedEdges)
    set { ← get with seats }
    pure { win with
      pointerMoveRequested := .none
      pointerResizeRequested := .none
    }
  set { ← get with windows }

def manageSeats (ctx : Context) : DisplayM State Unit := do
  let seats ← (← get).seats.mapM fun _ seat ↦ do
    for proxy in seat.interacted do
      let mut windowOrder := (← get).windowOrder
      let winIdx := windowOrder.findIdx (· == proxy)
      if h : winIdx < windowOrder.size then
        let win := windowOrder[winIdx]
        windowOrder := windowOrder.eraseIdx winIdx
        windowOrder := windowOrder.push win
      set { ← get with windowOrder }
    let mut seat := seat
    seat ← seat.focusTop
    seat ← seat.doAction ctx
    if seat.opRelease then
      seat ← seat.opEnd
    else
      seat.opManage
    pure { seat with
      interacted := .none
      opRelease := false
    }
  set { ← get with seats }

def handle (ctx : Context) : RiverWindowManagerV1.Event → DisplayM State Unit
  | .unavailable => do
      DisplayM.liftIO <| IO.println "Another window manager is already running"
      DisplayM.requestQuit
  | .finished => do
      DisplayM.requestQuit
  | .manageStart => do
      removeOutputs
      removeWindows
      removeSeats
      initNewWindows
      initNewSeats ctx
      manageWindows
      manageSeats ctx
      ctx.wm.manageFinish
  | .renderStart => do
      for seat in (← get).seats.values do
        match seat.op with
          | .none => pure ()
          | .move proxy startX startY =>
              DisplayM.zoom (mk_lens windows ≫ Std.HashMap.lens proxy) do
                Window.setPosition (startX + seat.opDeltaX) (startY + seat.opDeltaY)
          | .resize proxy startX startY startWidth startHeight edges =>
              DisplayM.zoom (mk_lens windows ≫ Std.HashMap.lens proxy) do
                let mut x := startX
                let mut y := startY
                for w in ← get do
                  if edges.left then x := x + startWidth  - w.width
                  if edges.top  then y := y + startHeight - w.height
                Window.setPosition x y
      ctx.wm.renderFinish
  | .sessionLocked => pure ()
  | .sessionUnlocked => pure ()
  | .window proxy => do
      let windows := (← get).windows.insert proxy {
        proxy
        node := ← proxy.getNode
      }
      let windowOrder := (← get).windowOrder.push proxy
      set { ← get with windows, windowOrder }
      DisplayM.zoom (mk_lens windows ≫ Std.HashMap.lens proxy) do
        proxy.listen handleWindow
  | .output proxy => do
      let outputs := (← get).outputs.insert proxy { proxy }
      set { ← get with outputs }
      DisplayM.zoom (mk_lens outputs ≫ Std.HashMap.lens proxy) do
        proxy.listen handleOutput
  | .seat proxy => do
      let seats := (← get).seats.insert proxy { proxy }
      set { ← get with seats }
      DisplayM.zoom (mk_lens seats ≫ Std.HashMap.lens proxy) do
        proxy.listen handleSeat

def launch : DisplayM State Unit := do
  _ ← DisplayM.registerSignal .SIGINT do
    DisplayM.liftIO <| IO.println "Received SIGINT"
    DisplayM.requestQuit
  zoomRegistry Registry.init
  DisplayM.sync do
    let wm : RiverWindowManagerV1 ← zoomRegistry Registry.bindSingleton
    let xkbBindings : RiverXkbBindingsV1 ← zoomRegistry Registry.bindSingleton
    wm.listen (handle { wm, xkbBindings })

def main : IO Unit :=
  DisplayM.run launch { }