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

open Wayland Protocols Waylean

def Array.findAndModifyM {α : Type} {m : Type → Type} [Monad m] (a : Array α) (pred : α → Bool) (update : α → m α) : m (Array α × Option α) := do
  let idx := a.findIdx pred
  if h : idx < a.size then
    let new ← update a[idx]
    return (a.setIfInBounds idx new, new)
  else return (a, .none)

def Array.findAndModify {α : Type} (a : Array α) (pred : α → Bool) (update : α → α) : Array α × Option α :=
  (a.findAndModifyM pred (update : α → Id α) : Id (Array α × Option α))

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.lens {α β : Type} [BEq α] [Hashable α] (key : α) : Lens (Std.HashMap α β) (Option β) where
  get x := x.get? key
  set v x := x.alter key fun _ ↦ v

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 : Array Window := #[]
  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 (proxy : RiverWindowV1) (ev : RiverWindowV1.Event) : DisplayM State Unit := do
  let windows := (← get).windows.findAndModify (fun win ↦ win.proxy == proxy) fun win ↦
    match ev with
    | .closed                     => { win with closed := true }
    | .dimensions w h             => { win with width := w, height := h }
    | .pointerMoveRequested s     => { win with pointerMoveRequested := s }
    | .pointerResizeRequested s e => { win with pointerResizeRequested := s, pointerResizeRequestedEdges := e }
    | _                           => win
  set { ← get with windows := windows.fst }

structure Context where
  wm : RiverWindowManagerV1
  xkbBindings : RiverXkbBindingsV1

def Window.setPosition {σ : Type} (w : Window) (x y : Int32) : DisplayM σ Window := do
  w.node.setPosition x y
  return { 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).windows
  if h : windows.size == 0 then
    s.proxy.clearFocus
    return { s with focused := .none }
  else
    let win := windows[windows.size - 1]'(by
      apply Nat.sub_one_lt
      rw [beq_iff_eq] at h
      assumption)
    s.proxy.focusWindow win.proxy
    win.node.placeTop
    return { s with focused := win.proxy }

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 windows := (← get).windows
        if h : 0 < windows.size then
          let win := windows[0]
          windows := windows.eraseIdx 0
          windows := windows.push win
        set { ← get with windows }
        s ← focusTop s
    | .move =>
        for proxy in s.hovered do
          for win in (← get).windows.find? (fun win ↦ win.proxy == proxy) do
            s ← s.pointerMove win
    | .resize =>
        for proxy in s.hovered do
          for win in (← get).windows.find? (fun win ↦ win.proxy == proxy) 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 Window.closed
  let mut seats := (← get).seats
  for win in closedWindows do
    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, 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
  let windows ← (← get).windows.mapM fun win ↦ do
    let mut win := win
    if win.new then
      win ← win.setPosition win.x win.y
      win.proxy.proposeDimensions win.width win.height
      win := { win with new := false }
    pure win
  set { ← get with windows }

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) do
      DisplayM.zoom (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 windows := (← get).windows
      let winIdx := windows.findIdx (fun win ↦ win.proxy == proxy)
      if h : winIdx < windows.size then
        let win := windows[winIdx]
        windows := windows.eraseIdx winIdx
        windows := windows.push win
      set { ← get with windows }
    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 =>
              let windows ← (← get).windows.findAndModifyM
                (fun win ↦ win.proxy == proxy)
                (fun win ↦ win.setPosition (startX + seat.opDeltaX) (startY + seat.opDeltaY))
              set { ← get with windows := windows.fst }
          | .resize proxy startX startY startWidth startHeight edges =>
              let windows ← (← get).windows.findAndModifyM
                (fun win ↦ win.proxy == proxy)
                (fun win ↦
                  let x := startX + if edges.left then (startWidth - win.width) else 0
                  let y := startY + if edges.top then (startHeight - win.height) else 0
                  win.setPosition x y)
              set { ← get with windows := windows.fst }
      ctx.wm.renderFinish
  | .sessionLocked => pure ()
  | .sessionUnlocked => pure ()
  | .window proxy => do
      let windows := (← get).windows.push {
        proxy
        node := ← proxy.getNode
      }
      set { ← get with windows }
      proxy.listen (handleWindow proxy)
  | .output proxy => do
      let outputs := (← get).outputs.insert proxy { proxy }
      set { ← get with outputs }
      DisplayM.zoom (mk_lens outputs) do
        DisplayM.zoom (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) do
        DisplayM.zoom (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 { }