local ui = {}

---@alias UiHandler fun(id: NodeId, x?: integer, y?: integer)

---@class UiSpec
---@field [integer] NodeId Child nodes.
---@field w? number|"fill"|"auto"
---@field h? number|"fill"|"auto"
---@field pad? number
---@field gap? number
---@field align? "start"|"center"|"end"|"stretch"
---@field justify? "start"|"center"|"end"|"between"
---@field row? boolean
---@field at? table
---@field label? string
---@field font? GuiFont
---@field style? GuiTextStyle
---@field fit? integer Maximum label width.
---@field background? GuiColor
---@field face? GuiColor
---@field pressedFace? GuiColor
---@field pressedColor? GuiColor
---@field focusColor? GuiColor
---@field focusWidth? integer
---@field textStyle? GuiTextStyle
---@field on_enter? UiHandler
---@field on_exit? UiHandler
---@field on_click? UiHandler
---@field paint? fun(id: NodeId, x: integer, y: integer, w: integer, h: integer, clipX: integer, clipY: integer, clipW: integer, clipH: integer)
---@field press_style? boolean False for a widget that paints its own press feedback.
---@field scrollX? boolean Children measure unbounded across, and the box pans horizontally.
---@field scrollY? boolean Children measure unbounded down, and the box pans vertically.

---@class UiConfirmSpec
---@field title string
---@field message? string
---@field ok? string
---@field cancel? string|false
---@field w? number
---@field background? GuiColor
---@field border? GuiColor
---@field on_ok? UiHandler
---@field on_cancel? UiHandler
---@field on_outside? UiHandler

local THEMES = {
  light = { background = { 255, 255, 255 }, color = { 0, 0, 0 }, accent = { 0, 120, 255 }, radius = 6 },
  dark = { background = { 18, 18, 20 }, color = { 235, 235, 235 }, accent = { 166, 118, 255 }, radius = 6 },
  mono = { background = { 255, 255, 255 }, color = { 0, 0, 0 }, accent = { 0, 0, 0 }, radius = 0 },
}

local enterHandlers, exitHandlers, clickHandlers, painters = {}, {}, {}, {}
-- A widget that paints its own press feedback opts out, so pressing one key does not
-- repaint the whole node the way a pressed style would.
local pressStyles = {}
local laidOut = false
local themeName
-- One tree per state, built by the function ui.mount() was given. Rebuilding is
-- cheap enough that nothing is retained between screens.
local builder, root, responder, insideResponder, confirming
local inset = 0
local applyPalette

-- Scroll gestures live in the tree, not the app: a drag on any scrollX/scrollY box pans it,
-- a flick coasts, and a child button's tap is suppressed once the drag passes SLOP. C++ owns
-- only the clamped offset (setScroll/getScroll); the momentum and the tap-vs-pan decision are
-- here, keyed by node so every scrollable box gets them for free.
local SCROLL_SLOP = 8
local SCROLL_DECAY = 0.02
local SCROLL_STOP = 8
-- id -> {x=bool, y=bool}: which axes a box pans. Recorded at build, cleared on reset.
local scrollNodes = {}
-- id -> {sx, sy, vx, vy}: pan offset (float, floored into setScroll) and flick velocity.
local scrollState = {}
-- Set of ids still coasting, iterated each draw until they settle.
local flinging = {}
-- Per-gesture bookkeeping: the box being panned once past SLOP, and the raw drag deltas.
local panning, scrollAncestor, dragged
local downX, downY, lastX, lastY, pendingX, pendingY

local function mix(a, b, amount)
  local out = {}
  for i = 1, 3 do
    out[i] = math.floor(a[i] + (b[i] - a[i]) * amount + 0.5)
  end
  return out
end

local function color(rgb)
  return screen.color(rgb[1], rgb[2], rgb[3])
end

local function palette(seed)
  local background, foreground, accent = seed.background, seed.color, seed.accent
  return {
    background = color(background),
    color = color(foreground),
    muted = color(mix(foreground, background, 0.45)),
    accent = color(accent),
    disabled = color(mix(background, foreground, 0.25)),
    face = color(mix(background, foreground, 0.08)),
    pressedFace = color(accent),
    pressedColor = color(background),
    focusColor = color(accent),
    radius = seed.radius,
  }
end

local function loadTheme(name)
  themeName = THEMES[name] and name or "light"
  ui.theme = palette(THEMES[themeName])
end

---@return string
function ui.getTheme()
  return themeName
end

---@return string[]
function ui.themeNames()
  local names = {}
  for name in pairs(THEMES) do
    names[#names + 1] = name
  end
  table.sort(names)
  return names
end

---@param name string
---@return true? ok
---@return string? error
function ui.setTheme(name)
  if not THEMES[name] then
    return nil, "Unknown theme"
  end
  local ok, err = screen.setTheme(name)
  if not ok then
    return nil, err
  end
  loadTheme(name)
  if root then
    applyPalette(root)
    screen.clear(ui.theme.background)
    tree.invalidate(root)
  end
  return true
end

loadTheme(screen.getTheme())

local function clearState()
  enterHandlers, exitHandlers, clickHandlers, painters = {}, {}, {}, {}
  pressStyles = {}
end

-- The clip is the slice of the node being repainted now, which for a composited repaint is
-- one band. A painter that ignores it still draws correctly; one that culls to it stops
-- redrawing its whole contents once per band it spans.
tree.setPainter(function(id, x, y, w, h, clipX, clipY, clipW, clipH)
  local painter = painters[id]
  if painter then
    painter(id, x, y, w, h, clipX, clipY, clipW, clipH)
  end
end)

local STYLE_KEYS = {
  "color",
  "fill",
  "border",
  "face",
  "pressedFace",
  "pressedColor",
  "focusColor",
  "radius",
  "font",
  "textStyle",
}

local function applyStyle(id, spec)
  local style, hasStyle = {}, false
  for _, key in ipairs(STYLE_KEYS) do
    if spec[key] ~= nil then
      style[key], hasStyle = spec[key], true
    end
  end
  if spec.background ~= nil then
    style.background, style.fill, hasStyle = spec.background, spec.background, true
  end
  if hasStyle then
    tree.setStyle(id, style)
  end
end

local function build(spec, kind)
  spec = spec or {}
  -- Nodes outside a build would reset the arena under the screen already on the
  -- panel, chrome included. Rebuilding is the only way to change one.
  if laidOut then
    error("build nodes from the function ui.mount() was given, then ui.rebuild()", 3)
  end

  local children = {}
  for index, child in ipairs(spec) do
    children[index] = child
    spec[index] = nil
  end

  spec.type = kind
  spec.interactive = spec.on_enter ~= nil or spec.on_exit ~= nil or spec.on_click ~= nil
  local id = tree.create(nil, spec)
  for _, child in ipairs(children) do
    tree.attach(id, child)
  end

  applyStyle(id, spec)
  enterHandlers[id] = spec.on_enter
  exitHandlers[id] = spec.on_exit
  clickHandlers[id] = spec.on_click
  painters[id] = spec.paint
  pressStyles[id] = spec.press_style ~= false
  if spec.scrollX or spec.scrollY then
    scrollNodes[id] = { x = spec.scrollX == true, y = spec.scrollY == true }
    scrollState[id] = { sx = 0, sy = 0, vx = 0, vy = 0 }
  end
  return id
end

---@param spec UiSpec
---@return NodeId
function ui.box(spec)
  return build(spec, "box")
end

---Side length for a grid of square cards that fits the frame, and the column count that
---produced it. Landscape gets a wider grid, so the same screen reflows on rotation.
---@param count integer Cards to place.
---@param pad integer Padding outside the grid.
---@param gap integer Gap between cards.
---@param reserve? integer Height to leave free below the grid.
---@return integer side
---@return integer columns
function ui.cardSide(count, pad, gap, reserve)
  local width, height = ui.frame()
  local columns = width >= height and 3 or 2
  local rows = math.ceil(count / columns)
  local byWidth = (width - 2 * pad - (columns - 1) * gap) // columns
  local byHeight = (height - 2 * pad - (reserve or 0) - (rows - 1) * gap) // rows
  return math.min(byWidth, byHeight), columns
end

---How much of the panel chrome took before the app was built. Set by whatever mounts the
---tree, because layout has not run yet when an app sizes itself.
---@param px integer
function ui.setInset(px)
  inset = px
end

---The box the app is built into, which is the panel minus the chrome above it.
---@return integer w
---@return integer h
function ui.frame()
  return screen.getWidth(), screen.getHeight() - inset
end

---@param spec UiSpec
---@return NodeId
function ui.spacer(spec)
  spec = spec or {}
  return build({ w = spec.w, h = spec.h }, "box")
end

---@param text string
---@param spec? UiSpec
---@return NodeId
function ui.text(text, spec)
  spec = spec or {}
  spec.label = text
  spec.textStyle, spec.style = spec.style, nil
  return build(spec, "text")
end

---@param text string
---@param spec? UiSpec
---@return NodeId
function ui.label(text, spec)
  spec = spec or {}
  local font, style = spec.font or screen.FONT_UI, spec.style or screen.STYLE_NORMAL
  if spec.fit and screen.getTextWidth(font, text, style) > spec.fit then
    while #text > 1 and screen.getTextWidth(font, text .. "~", style) > spec.fit do
      text = text:sub(1, -2)
    end
    text = text .. "~"
  end
  spec.w = screen.getTextWidth(font, text, style)
  spec.h = screen.getFontHeight(font, style)
  spec.font, spec.fit = font, nil
  return ui.text(text, spec)
end

---@param spec UiSpec
---@return NodeId
function ui.button(spec)
  spec = spec or {}
  spec.pad = spec.pad or 8
  spec.align = spec.align or "center"
  local label, font = spec.label, spec.font
  spec.label = nil
  local id = build(spec, "button")
  if label then
    tree.create(id, { type = "text", label = label, font = font or screen.FONT_UI })
  end
  return id
end

---@param spec UiSpec
---@return NodeId
function ui.custom(spec)
  return build(spec, "custom")
end

---@param id NodeId
---@param text string
function ui.setText(id, text)
  if tree.getLabel(id) == text then
    return
  end
  tree.setLabel(id, text)
  tree.invalidate(id)
end

---@param id NodeId
function ui.invalidate(id)
  tree.invalidate(id)
end

---Pans a scrollable node, clamped to its content. The node is marked dirty, so the next
---ui.draw() repaints it -- there is nothing for the app to draw and no handle to refresh,
---which is the whole point of scrolling in the tree rather than in a painter.
---@param id NodeId
---@param x integer
---@param y integer
function ui.setScroll(id, x, y)
  tree.setScroll(id, x, y)
end

---@param id NodeId
---@return integer x, integer y
function ui.getScroll(id)
  return tree.getScroll(id)
end

---How far each axis can pan. Zero on an axis whose content already fits.
---@param id NodeId
---@return integer maxX, integer maxY
function ui.getScrollRange(id)
  return tree.getScrollRange(id)
end

---@param spec UiConfirmSpec
---@return NodeId
function ui.confirm(spec)
  local card = {
    w = spec.w or 0.85,
    pad = 16,
    gap = 12,
    background = spec.background or ui.theme.background,
    border = spec.border or ui.theme.muted,
    ui.text(spec.title),
  }
  if spec.message then
    card[#card + 1] = ui.text(spec.message, { color = ui.theme.muted })
  end

  local buttons = { row = true, gap = 8, justify = "end" }
  if spec.cancel ~= false then
    buttons[#buttons + 1] = ui.button { label = spec.cancel or "cancel", on_click = spec.on_cancel }
  end
  buttons[#buttons + 1] = ui.button { label = spec.ok or "ok", on_click = spec.on_ok }
  card[#card + 1] = ui.box(buttons)

  return ui.box {
    at = { x = 0, y = 0 },
    w = "fill",
    h = "fill",
    -- A full-bleed layer painted last: hit-testing returns the topmost node, so taps on the
    -- dim area land here and cannot reach the content beneath. on_outside handles them; with
    -- none, they die at the root -- which is what makes the dialog modal.
    align = "center",
    justify = "center",
    on_click = spec.on_outside,
    ui.box(card),
  }
end

function ui.reset()
  tree.reset()
  clearState()
  laidOut = false
  root, responder, insideResponder, confirming = nil, nil, nil, nil
  scrollNodes, scrollState, flinging = {}, {}, {}
  panning, scrollAncestor, dragged = nil, nil, false
end

applyPalette = function(root)
  -- The root is the panel background, not a card: no border, so it takes the fast fillRect
  -- path rather than the per-pixel roundRect one. Radius stays so cards inherit it.
  tree.setStyle(root, {
    color = ui.theme.color,
    background = ui.theme.background,
    face = ui.theme.face,
    pressedFace = ui.theme.pressedFace,
    pressedColor = ui.theme.pressedColor,
    focusColor = ui.theme.focusColor,
    radius = ui.theme.radius,
    font = screen.FONT_UI,
  })
end

---Registers the function that builds the whole tree and shows what it returns.
---@param fn fun(): NodeId
function ui.mount(fn)
  builder = fn
  ui.rebuild()
end

---Rebuilds the tree from scratch and repaints. Screens are not retained, so this
---is how a screen changes, a rotation is answered and a dialog opens.
function ui.rebuild()
  ui.reset()
  root = builder()
  tree.setSize(root, "fill", "fill")
  applyPalette(root)
  local ok, err = tree.layout(root, 0, 0, screen.getWidth(), screen.getHeight())
  if not ok then
    error(err, 2)
  end
  tree.dropScratch()
  laidOut = true
  screen.clear(ui.theme.background)
  tree.draw(root)
  -- A build allocates a spec table per node and drops them all here, and the next thing an
  -- app does may be the one that needs a contiguous WiFi buffer. Collecting now costs a few
  -- milliseconds on a screen change nobody can see, and leaves the heap in a known state
  -- instead of one that depends on when the incremental GC last ran.
  collectgarbage()
end

-- Pushes a node's float offset into the tree, then snaps our copy back to whatever the
-- clamp allowed -- without this the offset runs past the end while the content sits still
-- and the list ignores the first part of the drag back.
local function applyScroll(id, st)
  ui.setScroll(id, math.floor(st.sx), math.floor(st.sy))
  local cx, cy = ui.getScroll(id)
  if cx and cx ~= math.floor(st.sx) then
    st.sx, st.vx = cx, 0
  end
  if cy and cy ~= math.floor(st.sy) then
    st.sy, st.vy = cy, 0
  end
end

-- Free settle: velocity decays to rest, applied as DECAY^seconds so the glide lasts the
-- same wall time whatever the frame rate. A paged box would instead ease toward the nearest
-- page boundary here -- same gesture, different ending -- so this is the one seam paging adds.
local function settleFree(st, seconds)
  st.sx, st.sy = st.sx + st.vx * seconds, st.sy + st.vy * seconds
  local keep = SCROLL_DECAY ^ seconds
  st.vx, st.vy = st.vx * keep, st.vy * keep
  if math.abs(st.vx) < SCROLL_STOP then
    st.vx = 0
  end
  if math.abs(st.vy) < SCROLL_STOP then
    st.vy = 0
  end
end

-- One GC step a frame, because the collector's default pace is the painter's problem: a
-- band buffer needs a contiguous block, and letting the heap double before a cycle lets
-- garbage take the block the band was going to get. Stepping keeps the sawtooth shallow
-- enough that beginBuffer() keeps succeeding instead of falling back to the panel.
---@param deltaMs? integer Elapsed frame time, used to advance any active pan or flick.
function ui.draw(deltaMs)
  local seconds = (deltaMs or 0) / 1000
  if seconds > 0 then
    if panning then
      local st = scrollState[panning]
      st.sx, st.sy = st.sx + pendingX, st.sy + pendingY
      -- Velocity is measured over the frame the motion arrived in, so a finger that paused
      -- before lifting reports zero and does not flick.
      st.vx, st.vy = pendingX / seconds, pendingY / seconds
      pendingX, pendingY = 0, 0
      applyScroll(panning, st)
    end
    for id in pairs(flinging) do
      local st = scrollState[id]
      settleFree(st, seconds)
      applyScroll(id, st)
      if st.vx == 0 and st.vy == 0 then
        flinging[id] = nil
      end
    end
  end
  if root then
    tree.draw(root)
  end
  collectgarbage "step"
end

-- The nearest ancestor (or self) that pans, or nil. A drag on a child button scrolls the
-- list it sits in, which is why the tap has to yield to the pan rather than the reverse.
local function scrollableAncestor(id)
  while id do
    if scrollNodes[id] then
      return id
    end
    id = tree.getParent(id)
  end
end

-- The nearest ancestor (or self) that handles a press, or nil. tree.hit returns the deepest
-- node by geometry -- a button's text label, say -- so dispatch bubbles up to whoever owns
-- the on_click/on_enter/on_exit, the way a responder chain does.
local function handlerAncestor(id)
  while id do
    if clickHandlers[id] or enterHandlers[id] or exitHandlers[id] then
      return id
    end
    id = tree.getParent(id)
  end
end

local function inside(id, x, y)
  local rx, ry, rw, rh = tree.getRect(id)
  return x >= rx and x < rx + rw and y >= ry and y < ry + rh
end

local function enter(id, x, y)
  if pressStyles[id] then
    tree.setPressed(id, true)
  end
  local handler = enterHandlers[id]
  if handler then
    handler(id, x, y)
  end
end

local function exit(id, x, y)
  if pressStyles[id] then
    tree.setPressed(id, false)
  end
  local handler = exitHandlers[id]
  if handler then
    handler(id, x, y)
  end
end

---@param x integer
---@param y integer
---@return boolean handled
function ui.down(x, y)
  local focused = tree.getFocus()
  if focused then
    tree.setFocus(nil)
    local handler = exitHandlers[focused]
    if handler then
      handler(focused)
    end
  end

  local hitNode = root and tree.hit(root, x, y)
  if not hitNode then
    return false
  end
  responder, scrollAncestor = handlerAncestor(hitNode), scrollableAncestor(hitNode)
  insideResponder, dragged, panning = responder ~= nil, false, nil
  downX, downY, lastX, lastY = x, y, x, y
  pendingX, pendingY = 0, 0
  if scrollAncestor then
    -- A finger down catches an in-progress glide, the way every touch UI does.
    flinging[scrollAncestor] = nil
    local st = scrollState[scrollAncestor]
    st.vx, st.vy = 0, 0
  end
  if responder then
    enter(responder, x, y)
  end
  return responder ~= nil or scrollAncestor ~= nil
end

---@param x integer
---@param y integer
---@return boolean handled
function ui.move(x, y)
  if not responder and not scrollAncestor then
    return false
  end
  if scrollAncestor then
    local flags = scrollNodes[scrollAncestor]
    if flags.x then
      pendingX = pendingX + (lastX - x)
    end
    if flags.y then
      pendingY = pendingY + (lastY - y)
    end
    lastX, lastY = x, y
    if not dragged then
      local past = (flags.x and math.abs(x - downX) > SCROLL_SLOP)
        or (flags.y and math.abs(y - downY) > SCROLL_SLOP)
      if past then
        dragged, panning = true, scrollAncestor
        -- The press became a scroll: drop the responder's feedback and keep it from clicking.
        if responder and insideResponder then
          exit(responder, x, y)
        end
        insideResponder = false
      end
    end
    if dragged then
      return true
    end
  end
  if responder then
    local isInside = inside(responder, x, y)
    if isInside ~= insideResponder then
      insideResponder = isInside
      if isInside then
        enter(responder, x, y)
      else
        exit(responder, x, y)
      end
    end
  end
  return true
end

---@param x integer
---@param y integer
---@return boolean handled
function ui.up(x, y)
  if not responder and not scrollAncestor then
    return false
  end
  if dragged then
    -- The gesture was a pan: hand any velocity the last frames built to draw(), which coasts
    -- and clamps it. A drag that ended still carries zero velocity, so it simply stops.
    if panning then
      flinging[panning] = true
    end
    responder, insideResponder, scrollAncestor, dragged, panning = nil, nil, nil, false, nil
    return true
  end
  local target, wasActive = responder, insideResponder
  responder, insideResponder, scrollAncestor = nil, nil, nil
  if target then
    if wasActive then
      exit(target, x, y)
    end
    if inside(target, x, y) then
      local handler = clickHandlers[target]
      if handler then
        handler(target, x, y)
      end
    end
  end
  return true
end

local DIRECTIONS = { up = true, down = true, left = true, right = true }

local function focusFirst()
  local focused = tree.focusFirst(root)
  if focused then
    local handler = enterHandlers[focused]
    if handler then
      handler(focused)
    end
  end
  return focused
end

---@param name string Button name; directions and confirm are handled.
---@param pressed boolean
---@return boolean handled
function ui.buttonPress(name, pressed)
  if type(pressed) ~= "boolean" then
    error("button state must be boolean", 2)
  end

  if DIRECTIONS[name] then
    if not pressed then
      return true
    end
    local previous = tree.getFocus()
    if not previous then
      focusFirst()
      return true
    end
    local focused = tree.moveFocus(root, name)
    if focused ~= previous then
      local leave = exitHandlers[previous]
      if leave then
        leave(previous)
      end
      local arrive = enterHandlers[focused]
      if arrive then
        arrive(focused)
      end
    end
    return true
  end

  if name ~= "confirm" then
    return false
  end
  local focused = tree.getFocus() or focusFirst()
  if not focused then
    return false
  end
  if pressed then
    tree.setPressed(focused, true)
    confirming = focused
  else
    local target = confirming
    confirming = nil
    if target then
      tree.setPressed(target, false)
      local handler = clickHandlers[target]
      if handler then
        handler(target)
      end
    end
  end
  return true
end

return ui