-- Columns: a custom (Lua) tiling layout — a row of vertical columns. Think of it as
-- one long horizontal strip of columns that spans every monitor: focus/move flow off
-- the edge of one monitor onto the next, so it behaves like a huge row of scrolling
-- workspaces rather than a set of isolated boxes.
--
-- Each column is an INDEPENDENTLY-SCROLLING vertical strip (the whole point): when a
-- column holds more windows than fit, it pans to keep the focused one in view while the
-- other columns stay put — like the wheels of a combination lock. Sizing is flexible,
-- nothing is a fixed fraction:
--   * Each COLUMN has a width weight; the viewport width is split between columns by
--     those weights. Resizing left/right grows/shrinks the focused column.
--   * Each WINDOW has a height weight. The base height of a weight-1 window is
--     viewport/`rows` (rows = windows-per-view). A column that is *underfull* stretches
--     its windows to fill (a lone window fills the column; two equal windows share it
--     50/50) — so you can stack a file-explorer over a terminal in column 1 while
--     Firefox owns column 2. A column that is *overfull* keeps natural heights and
--     SCROLLS. Resizing up/down retunes the focused window's height weight.
-- Weights and per-column pan persist across relayouts.
--
-- This is a real custom layout (hl.layout.register), not the built-in scrolling engine —
-- that one has a single shared tape; this pans each column apart.
--
-- Controls (routed here from binds.lua via `hl.dsp.layout("<cmd>")` → layout_msg):
--   focus l/r/u/d   — l/r step between columns (off the edge → next/prev monitor),
--                     u/d step within the focused column (which pans to it). Swap vertical.
--   move  l/r/u/d   — l/r move between columns (off the edge → next/prev monitor).
--                     u/d PARK/UNPARK within the column: moving toward a single neighbour
--                     parks the two side by side in one slot (persistent; both keep the full
--                     scroll-axis length, split across the column — a vertical divider when
--                     horizontal — and the column closes up so there's no gap). Moving a
--                     parked window again pops it out and past its partner.
--   resize l/r/u/d  — l/r change the focused column's width, u/d the focused window's
--                     height (weights, so neighbours give/take space). Axes swap vertical.
--   count h/j/k/l   — horizontal ws: h/l = −/＋ columns, k/j = −/＋ windows-per-view.
--                     vertical ws:   k/j = −/＋ columns, h/l = −/＋ per-view.
--   cols +/-/N      — direct column-count set (the menu's stepper).
--   orient h/v      — orientation (set from the menu dropdown).
--   center on/off/toggle — "center-focused" switch (menu): when a column is overfull
--                     and scrolls, keep the focused window centred in the column
--                     instead of the minimal-movement default (whichever edge of the
--                     focused window would otherwise go off-screen just barely comes
--                     back into view). Global (mirrors the scrolling layout's
--                     "Centered follow" toggle), persisted across reloads, OFF by
--                     default. Only ever repositions the column holding the focused
--                     window — every other column's pan is untouched, since there is
--                     only one focused window at a time.
--
-- State is per workspace; a window's column assignment, order, both weight tables and
-- each column's pan offset persist across relayouts.

local state = {}   -- [wsid] = { ncols, rows, orient, assign={addr=col}, seq={addr=n},
                   --            nextseq, colw={col=weight}, winh={addr=weight}, pan={col=px} }

-- Resize weights: base 1.0, clamped to a sane band so nothing collapses or hogs.
local RSTEP, RMIN, RMAX = 0.15, 0.25, 6.0
local function clampw(v) return math.max(RMIN, math.min(RMAX, v)) end

-- Publish per-workspace column count so the astal-menu can show a live stepper.
local STATE_FILE = (os.getenv("HOME") or "") .. "/.cache/astal-menu/columns-state.json"

-- "center-focused" switch — global (like the scrolling layout's own fit-method
-- toggle), so a single state, not per-workspace. Seeded once at module load from
-- the last-published state, since a reload wipes this module's Lua state (same
-- reason ncols/rows/orient are seeded — see seed() below).
local center_focus = false
do
    local f = io.open(STATE_FILE, "r")
    if f then
        local data = f:read("*a") or ""; f:close()
        if data:match('"_center"%s*:%s*true') then center_focus = true end
    end
end

local function publish()
    local parts = { string.format('"_center":%s', tostring(center_focus)) }
    for wsid, s in pairs(state) do
        parts[#parts + 1] = string.format('"%s":{"cols":%d,"rows":%d,"orient":"%s"}',
            tostring(wsid), s.ncols, s.rows, s.orient)
    end
    pcall(function()
        local f = io.open(STATE_FILE, "w")
        if f then f:write("{" .. table.concat(parts, ",") .. "}\n"); f:close() end
    end)
end

-- Seed ncols/rows/orient for a workspace from the last published state — a reload wipes
-- this module's Lua state, so without this the user's column count and orientation reset
-- to defaults every reload. (Per-window assignment/weights are not restored; they reflow.)
local function seed(wsid)
    local out = {}
    local f = io.open(STATE_FILE, "r")
    if not f then return out end
    local data = f:read("*a") or ""; f:close()
    local blob = data:match('"' .. tostring(wsid) .. '"%s*:%s*(%b{})')
    if not blob then return out end
    out.ncols  = tonumber(blob:match('"cols"%s*:%s*(%d+)'))
    out.rows   = tonumber(blob:match('"rows"%s*:%s*(%d+)'))
    out.orient = blob:match('"orient"%s*:%s*"(%a+)"')
    return out
end

-- First-ever layout of a workspace (no seeded/published state) picks its starting
-- column count from the monitor's usable width, ~COL_TARGET_WIDTH px per column, so a
-- wide monitor doesn't start out with one cramped default column and a narrow one
-- doesn't start out overcrowded. 1080p (1920px) / 700 -> 2 columns.
local COL_TARGET_WIDTH = 700

local function ws_state(wsid, area)
    local s = state[wsid]
    if not s then
        local sd = seed(wsid)
        local orient = sd.orient or "h"
        local ncols = sd.ncols
        if not ncols then
            local len = area and ((orient == "h") and area.w or area.h)
            ncols = (len and len > 0) and math.max(1, math.floor(len / COL_TARGET_WIDTH)) or 2
        end
        s = { ncols = ncols, rows = sd.rows or 2, orient = orient,
              assign = {}, seq = {}, nextseq = 1, colw = {}, winh = {}, pan = {}, pair = {} }
        state[wsid] = s
        publish()
    end
    return s
end

-- Reflow every window evenly across the current column count, by insertion order.
local function reflow(s, targets)
    local addrs = {}
    for _, tg in ipairs(targets) do
        local ok, addr = pcall(function() return tg.window.address end)
        if ok and addr then addrs[#addrs + 1] = addr end
    end
    table.sort(addrs, function(a, b) return (s.seq[a] or 0) < (s.seq[b] or 0) end)
    for i, addr in ipairs(addrs) do s.assign[addr] = ((i - 1) % s.ncols) + 1 end
    for k in pairs(s.pair) do s.pair[k] = nil end   -- a redistribute breaks parked pairs
end

local function win_ws_id(win)
    local ok, ws = pcall(function() return win.workspace end)
    if ok and ws ~= nil then
        local ok2, id = pcall(function() return ws.id end)
        if ok2 and id then return id end
        local id2 = tostring(ws):match("%((%-?%d+)")
        if id2 then return tonumber(id2) end
    end
    return 0
end

local function focused_addr()
    local aw = hl.get_active_window and hl.get_active_window()
    if not aw then return nil end
    local ok, a = pcall(function() return aw.address end)
    return ok and a or nil
end

-- Build the per-column ordered address lists from the live targets, reconciling the
-- persistent assignment (drop closed windows; a single new window stacks into the
-- working column, a batch spreads across the emptiest columns — see below).
local function columnize(s, targets)
    local present, tgt_of = {}, {}
    for _, tg in ipairs(targets) do
        local ok, addr = pcall(function() return tg.window.address end)
        if ok and addr then present[addr] = true; tgt_of[addr] = tg end
    end
    for addr in pairs(s.assign) do
        if not present[addr] then
            local p = s.pair[addr]
            if p then s.pair[p] = nil; s.pair[addr] = nil end   -- release a parked partner
            s.assign[addr] = nil; s.seq[addr] = nil; s.winh[addr] = nil
        end
    end
    local foc = focused_addr()
    local counts = {}
    for c = 1, s.ncols do counts[c] = 0 end
    for addr, c in pairs(s.assign) do
        if present[addr] and c >= 1 and c <= s.ncols then counts[c] = counts[c] + 1 end
    end
    local fresh = {}
    for addr in pairs(present) do if not s.assign[addr] then fresh[#fresh + 1] = addr end end
    local function least_col()
        local best, bestn = 1, math.huge
        for c = 1, s.ncols do if counts[c] < bestn then best, bestn = c, counts[c] end end
        return best
    end
    if #fresh == 1 and next(s.assign) ~= nil then
        -- ONE new window while others already exist → it stacks into the column the user
        -- is working in (the new window itself grabbed focus, so use the remembered
        -- last-focused column). Opening a 2nd window thus shares that column instead of
        -- spreading to an empty one — that's how you get two windows at one height.
        local c = s.lastcol
        if not (c and c >= 1 and c <= s.ncols) then
            -- focus memory not ready (a just-opened window's focus can lag the relayout)
            -- → join the most-recently-opened window's column instead.
            local best_seq = -1
            for addr, col in pairs(s.assign) do
                if present[addr] and (s.seq[addr] or 0) > best_seq then
                    best_seq = s.seq[addr]; c = col
                end
            end
            if not (c and c >= 1 and c <= s.ncols) then c = least_col() end
        end
        -- ...but don't let one column run away: if it's already 2+ ahead of the emptiest,
        -- send the newcomer to the emptiest instead. This both balances rapid opens and
        -- keeps a reload (which re-adds every window one-by-one) from piling into column 1.
        local lc = least_col()
        if counts[c] - counts[lc] >= 2 then c = lc end
        s.assign[fresh[1]] = c; counts[c] = counts[c] + 1
        s.seq[fresh[1]] = s.nextseq; s.nextseq = s.nextseq + 1
    else
        -- first window, or several appearing at once (session start / restore) → spread
        -- across the emptiest columns so a fresh session isn't one giant pile.
        for _, addr in ipairs(fresh) do
            local c = least_col()
            s.assign[addr] = c; counts[c] = counts[c] + 1
            s.seq[addr] = s.nextseq; s.nextseq = s.nextseq + 1
        end
    end
    for addr, c in pairs(s.assign) do
        if c > s.ncols then s.assign[addr] = s.ncols elseif c < 1 then s.assign[addr] = 1 end
    end
    -- remember the focused window's column so the *next* new window stacks there.
    if foc and s.assign[foc] then s.lastcol = s.assign[foc] end
    local cols = {}
    for c = 1, s.ncols do cols[c] = {} end
    for addr in pairs(present) do table.insert(cols[s.assign[addr]], addr) end
    for c = 1, s.ncols do
        table.sort(cols[c], function(a, b) return (s.seq[a] or 0) < (s.seq[b] or 0) end)
    end
    return cols, tgt_of, foc
end

local function recalculate(ctx)
    local t = ctx.targets
    if type(t) ~= "table" or not t[1] then return end
    local A = ctx.area
    if type(A) ~= "table" then return end
    local s = ws_state(win_ws_id(t[1].window), A)
    local cols, tgt_of, foc = columnize(s, t)
    local horizontal = (s.orient ~= "v")
    local rows = math.max(1, s.rows)

    -- primary axis is split between columns; the secondary axis stacks each column's
    -- windows and pans (scrolls) when they overflow.
    local primLen = horizontal and A.w or A.h
    local primOrg = horizontal and A.x or A.y
    local secLen  = horizontal and A.h or A.w
    local secOrg  = horizontal and A.y or A.x

    local csum = 0
    for c = 1, s.ncols do csum = csum + (s.colw[c] or 1) end
    if csum <= 0 then csum = 1 end

    local cpos = primOrg
    for c = 1, s.ncols do
        local colLen = primLen * ((s.colw[c] or 1) / csum)
        local list = cols[c]
        local m = #list
        if m > 0 then
            -- group the column into rows: two mutually-paired, adjacent windows share one
            -- row side by side (a "parked" pair); every other window is its own row. Pairing
            -- removes a row, so the column closes up around a park — no gap.
            local rowlist = {}
            local i = 1
            while i <= m do
                local a, b = list[i], list[i + 1]
                if b and s.pair[a] == b and s.pair[b] == a then
                    rowlist[#rowlist + 1] = { a, b }; i = i + 2
                else
                    rowlist[#rowlist + 1] = { a }; i = i + 1
                end
            end

            -- base height of a weight-1 row is viewport / windows-per-view; a parked row
            -- takes the larger of its two weights so neither window is clipped.
            local unit = secLen / rows
            local function wof(addr) local w = s.winh[addr] or 1; return (w > 0) and w or 1 end
            local rh, total = {}, 0
            for r, row in ipairs(rowlist) do
                local w = (#row == 2) and math.max(wof(row[1]), wof(row[2])) or wof(row[1])
                rh[r] = w * unit; total = total + rh[r]
            end

            -- row holding the focused window (drives the pan)
            local frow
            for r, row in ipairs(rowlist) do
                for _, addr in ipairs(row) do if addr == foc then frow = r end end
            end

            local pan
            if total <= secLen then
                -- underfull: stretch rows to fill so there's no dead space.
                local scale = secLen / total
                for r = 1, #rowlist do rh[r] = rh[r] * scale end
                pan = 0
            else
                -- overfull: keep natural heights and PAN so the focused row stays in view
                -- (center_focus: centred in the viewport instead of minimal movement).
                -- Only the column holding the focused window (frow ~= nil) is touched here;
                -- every other column keeps whatever pan it already had.
                pan = s.pan[c] or 0
                if frow then
                    local top = 0
                    for r = 1, frow - 1 do top = top + rh[r] end
                    local bottom = top + rh[frow]
                    if center_focus then
                        pan = top + (rh[frow] - secLen) / 2
                    else
                        if top - pan < 0 then pan = top end
                        if bottom - pan > secLen then pan = bottom - secLen end
                    end
                end
                -- clamp to the column's actual scroll range either way — this is also what
                -- keeps "centred" from showing blank space past the first/last row.
                pan = math.max(0, math.min(pan, total - secLen))
            end
            s.pan[c] = pan

            -- place a box given position/length on the primary (column) axis and the
            -- secondary (scroll/stacking) axis, mapped to x/y per orientation.
            local function place_box(addr, ppos, plen, spos, slen)
                local box = horizontal
                    and { x = ppos, y = spos, w = plen, h = slen }
                    or  { x = spos, y = ppos, w = slen, h = plen }
                local tg = tgt_of[addr]
                if tg then pcall(function() tg:place(box) end) end
            end

            local spos = secOrg - pan
            for r, row in ipairs(rowlist) do
                if #row == 1 then
                    place_box(row[1], cpos, colLen, spos, rh[r])
                else
                    -- parked pair: side by side, split across the column, both keeping the
                    -- full row length along the scroll axis (a vertical divider horizontally).
                    local halfp = colLen / 2
                    place_box(row[1], cpos, halfp, spos, rh[r])
                    place_box(row[2], cpos + halfp, halfp, spos, rh[r])
                end
                spos = spos + rh[r]
            end
        end
        cpos = cpos + colLen
    end
end

local function layout_msg(ctx, cmd)
    local verb, arg = tostring(cmd or ""):match("^(%S+)%s*(%S*)$")
    if not verb then return end

    -- "center-focused" switch: global, doesn't touch any window, so handled before the
    -- empty-workspace bail-out below (toggling it from the menu shouldn't require a
    -- window to be focused/present).
    if verb == "center" then
        if arg == "on" then center_focus = true
        elseif arg == "off" then center_focus = false
        else center_focus = not center_focus end
        publish()
        return
    end

    local t = ctx.targets
    if type(t) ~= "table" or not t[1] then return end
    local s = ws_state(win_ws_id(t[1].window))
    local cols, _, foc = columnize(s, t)
    local horizontal = (s.orient ~= "v")

    local fc = (foc and s.assign[foc]) or 1
    local fi = 1
    if foc then for i, a in ipairs(cols[fc] or {}) do if a == foc then fi = i end end end

    -- l/r (horizontal) or u/d (vertical) run perpendicular to the column axis → these
    -- act "between columns"; the other pair acts within the focused column.
    local function between(a) return horizontal and (a == "l" or a == "r") or (not horizontal and (a == "u" or a == "d")) end
    local function decrease(a) return a == "l" or a == "u" end
    local function unpair(addr)
        local p = addr and s.pair[addr]
        if p then s.pair[p] = nil; s.pair[addr] = nil end
    end

    if verb == "focus" then
        if between(arg) then
            local nc = fc + (decrease(arg) and -1 or 1)
            if nc < 1 or nc > s.ncols then
                -- Off the edge: jump to the monitor spatially in that screen direction
                -- (arg is already "l"/"r"/"u"/"d", i.e. the edge we just fell off of),
                -- not Hyprland's own -1/+1 list-order cycling — see monitor-adjacent for
                -- why that distinction matters once a monitor has a transform.
                hl.exec_cmd("~/.config/scripts/monitor-adjacent focus " .. arg)
            elseif cols[nc] and #cols[nc] > 0 then
                local ni = math.min(fi, #cols[nc])
                hl.dispatch(hl.dsp.focus({ window = "address:" .. cols[nc][ni] }))
            end
        else
            -- within-column step WRAPS at the ends (both keyboard focus u/d and the
            -- mouse wheel, which forwards here too — see hypr-nav's scroll case).
            local list = cols[fc]
            if list and #list > 0 then
                local n = #list
                local step = decrease(arg) and -1 or 1
                local ni = ((fi - 1 + step) % n + n) % n + 1
                hl.dispatch(hl.dsp.focus({ window = "address:" .. list[ni] }))
            end
        end
    elseif verb == "move" then
        if not foc then return end
        if between(arg) then
            unpair(foc)                                                -- leaving the column un-parks it
            local nc = fc + (decrease(arg) and -1 or 1)
            if nc < 1 or nc > s.ncols then
                -- Off the edge: carry the window to the monitor spatially in that screen
                -- direction — see the matching comment in the "focus" branch above.
                hl.exec_cmd("~/.config/scripts/monitor-adjacent move " .. arg)
            else
                s.assign[foc] = nc
            end
        else
            -- PARK / UNPARK within the column. Moving toward a single neighbour parks the
            -- two side by side in one slot — a persistent state, and the column closes up so
            -- there's no gap. Moving a parked window again pops it out and past its partner.
            local list = cols[fc]
            local step = decrease(arg) and -1 or 1
            local partner = s.pair[foc]
            if partner then
                -- parked → un-park and move the focused window past its partner.
                s.pair[foc] = nil; s.pair[partner] = nil
                local a, b = s.seq[foc] or 0, s.seq[partner] or 0
                if (step > 0 and a < b) or (step < 0 and a > b) then
                    s.seq[foc], s.seq[partner] = b, a                  -- put foc on the far side
                end
            elseif list then
                local nb = list[fi + step]
                if nb and not s.pair[nb] then
                    s.pair[foc] = nb; s.pair[nb] = foc                 -- park side by side
                elseif nb then
                    -- neighbour is already a parked pair → hop the focused window over both.
                    local nb2 = s.pair[nb]
                    local n1, n2 = s.seq[nb] or 0, s.seq[nb2] or 0
                    s.seq[foc] = (step > 0) and (math.max(n1, n2) + 0.5) or (math.min(n1, n2) - 0.5)
                end
            end
        end
    elseif verb == "resize" then
        if between(arg) then
            -- grow/shrink the focused column's width weight
            s.colw[fc] = clampw((s.colw[fc] or 1) + (decrease(arg) and -RSTEP or RSTEP))
        elseif foc then
            -- grow/shrink the focused window's height weight within its column
            s.winh[foc] = clampw((s.winh[foc] or 1) + (decrease(arg) and -RSTEP or RSTEP))
        end
    elseif verb == "count" then
        -- horizontal ws: h/l change column count, k/j change windows-per-view.
        -- vertical ws:   k/j change column count, h/l change windows-per-view.
        local more_cols = horizontal and "l" or "j"
        local less_cols = horizontal and "h" or "k"
        local more_rows = horizontal and "j" or "l"
        local less_rows = horizontal and "k" or "h"
        local changed_cols = false
        if arg == more_cols then s.ncols = s.ncols + 1; changed_cols = true
        elseif arg == less_cols then s.ncols = math.max(1, s.ncols - 1); changed_cols = true
        elseif arg == more_rows then s.rows = s.rows + 1
        elseif arg == less_rows then s.rows = math.max(1, s.rows - 1) end
        if changed_cols then reflow(s, t) end
        publish()
    elseif verb == "cols" then
        -- direct column-count change (the menu's (-)[N](+) stepper).
        if arg == "+" or arg == "+1" then s.ncols = s.ncols + 1
        elseif arg == "-" or arg == "-1" then s.ncols = math.max(1, s.ncols - 1)
        else local n = tonumber(arg); if n then s.ncols = math.max(1, math.floor(n)) end end
        reflow(s, t)
        publish()
    elseif verb == "orient" then
        s.orient = (arg == "v") and "v" or "h"
        publish()
    end
end

-- Register the layout (guard against the double-register that a live reload can cause).
pcall(function()
    hl.layout.register("columns", { recalculate = recalculate, layout_msg = layout_msg })
end)

-- Spec consumed by the layout registry (init.lua) + the astal-menu dropdown.
return {
    name        = "columns",
    label       = "Columns",
    backend     = "lua:columns",   -- custom Lua layout, addressed with the lua: prefix
    lua         = true,            -- init.lua sends orientation via layout_msg, not scrolling:direction
    directional = true,           -- reuse the menu's direction dropdown for orientation
    dirs        = { "h", "v" },
    dir_labels  = { "Horizontal", "Vertical" },
    default_dir = "h",
    fit_method  = true,    -- menu shows the "center-focused" switch (routed to "center on/off")
    stepper     = true,   -- menu shows a "Columns: (-)[N](+)" count stepper
}