-- Columns: a custom (Lua) tiling layout — a row of columns, each an INDEPENDENTLY
-- panning vertical strip. Moving focus up/down inside a column pans only that column
-- to keep the focused window in view; the other columns stay frozen. Like the wheels
-- of a combination lock, every column scrolls on its own.
--
-- This is a real custom layout (hl.layout.register), not the built-in scrolling
-- engine — that one has a single shared tape offset and cannot pan columns apart.
--
-- Controls (routed here from binds.lua via `hl.dsp.layout("<cmd>")` → layout_msg):
--   focus l/r/u/d   — l/r step between columns, u/d step within the focused column
--                     (which then pans to reveal the focus). Axes swap when vertical.
--   move  l/r/u/d   — move the focused window between columns / reorder in its column.
--   count h/j/k/l   — horizontal ws: h/l = −/＋ columns, k/j = −/＋ windows-per-view.
--                     vertical ws:   k/j = −/＋ columns, h/l = −/＋ per-view.
--   orient h/v      — orientation (set from the menu dropdown).
--
-- State is per workspace; a window's column assignment and each column's pan offset
-- persist across relayouts.

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

local function ws_state(wsid)
    local s = state[wsid]
    if not s then
        s = { ncols = 2, rows = 2, orient = "h", assign = {}, seq = {}, nextseq = 1, pan = {} }
        state[wsid] = s
    end
    return s
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, assign new ones to the focused column).
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 s.assign[addr] = nil; s.seq[addr] = nil end
    end
    local foc = focused_addr()
    -- new windows fill the least-populated column so they spread out.
    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
    for addr in pairs(present) do
        if not s.assign[addr] then
            local best, bestn = 1, math.huge
            for c = 1, s.ncols do if counts[c] < bestn then best, bestn = c, counts[c] end end
            s.assign[addr] = best; counts[best] = counts[best] + 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
    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))
    local cols, tgt_of, foc = columnize(s, t)
    local horizontal = (s.orient ~= "v")
    local rows = math.max(1, s.rows)

    for c = 1, s.ncols do
        local list = cols[c]
        local m = #list
        if m > 0 then
            -- primary axis is split into ncols; the secondary axis stacks the windows,
            -- each sized to 1/rows of the viewport, and pans.
            local colSpan, colStart, winSec, secStart, secTotal
            if horizontal then
                colSpan, colStart = A.w / s.ncols, A.x + (c - 1) * (A.w / s.ncols)
                winSec, secStart, secTotal = A.h / rows, A.y, A.h
            else
                colSpan, colStart = A.h / s.ncols, A.y + (c - 1) * (A.h / s.ncols)
                winSec, secStart, secTotal = A.w / rows, A.x, A.w
            end
            -- pan so the focused window (if in this column) stays in view; else clamp.
            local fidx
            for i, addr in ipairs(list) do if addr == foc then fidx = i end end
            local pan = s.pan[c] or 0
            if fidx then
                local top = (fidx - 1) * winSec
                if top - pan < 0 then pan = top end
                if top + winSec - pan > secTotal then pan = top + winSec - secTotal end
            end
            pan = math.max(0, math.min(pan, math.max(0, m * winSec - secTotal)))
            s.pan[c] = pan
            for i, addr in ipairs(list) do
                local off = (i - 1) * winSec - pan
                local box = horizontal
                    and { x = colStart, y = secStart + off, w = colSpan, h = winSec }
                    or  { x = secStart + off, y = colStart, w = winSec, h = colSpan }
                local tg = tgt_of[addr]
                if tg then pcall(function() tg:place(box) end) end
            end
        end
    end
end

local function layout_msg(ctx, cmd)
    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 verb, arg = tostring(cmd or ""):match("^(%S+)%s*(%S*)$")
    if not verb then return end
    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 or u/d perpendicular to the column axis means "between columns"
    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

    if verb == "focus" then
        if between(arg) then
            local nc = fc + (decrease(arg) and -1 or 1)
            if nc >= 1 and nc <= s.ncols and 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
            local ni = fi + (decrease(arg) and -1 or 1)
            if cols[fc] and ni >= 1 and ni <= #cols[fc] then
                hl.dispatch(hl.dsp.focus({ window = "address:" .. cols[fc][ni] }))
            end
        end
    elseif verb == "move" then
        if not foc then return end
        if between(arg) then
            s.assign[foc] = math.max(1, math.min(s.ncols, fc + (decrease(arg) and -1 or 1)))
        else
            local list = cols[fc]
            local ni = fi + (decrease(arg) and -1 or 1)
            if list and ni >= 1 and ni <= #list then
                local other = list[ni]
                s.seq[foc], s.seq[other] = s.seq[other], s.seq[foc]
            end
        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 all windows evenly across the new column count (by insertion order)
            local addrs = {}
            for _, tg in ipairs(t) 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
        end
    elseif verb == "orient" then
        s.orient = (arg == "v") and "v" or "h"
    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  = false,
}