-- which-key.lua
-- ============================================================================
-- Displays available keybindings when a prefix key is pressed and the user
-- waits for a configurable delay. Similar to Emacs which-key package.
--
-- This package uses the editor's idle time tracking and extended echo area
-- to show hints without requiring any C++ code changes.
-- ============================================================================

local which_key = {}

-- Configuration
which_key.config = {
    delay_ms = 1000,       -- Delay before showing hints (milliseconds)
    max_hints = 20,        -- Max number of hints to show
    max_lines = 5,         -- Max lines in echo area
    separator = "  ",      -- Separator between hints on same line
    show_description = true, -- Show command descriptions
    enabled = true,        -- Enable which-key
}

-- State
which_key.last_prefix = nil
which_key.showing_hints = false

-- Parse a key sequence like "C-x C-f" into parts
local function parse_sequence(seq)
    local parts = {}
    for part in seq:gmatch("%S+") do
        table.insert(parts, part)
    end
    return parts
end

-- Check if sequence starts with prefix
local function starts_with(sequence, prefix)
    local seq_parts = parse_sequence(sequence)
    local prefix_parts = parse_sequence(prefix)

    if #seq_parts < #prefix_parts then
        return false
    end

    for i, part in ipairs(prefix_parts) do
        if seq_parts[i] ~= part then
            return false
        end
    end

    return true
end

-- Get the next key after the prefix
local function get_next_key(sequence, prefix)
    local seq_parts = parse_sequence(sequence)
    local prefix_parts = parse_sequence(prefix)

    if #seq_parts > #prefix_parts then
        return seq_parts[#prefix_parts + 1]
    end

    return nil
end

-- Get all bindings that start with the given prefix
function which_key.get_bindings_for_prefix(prefix)
    local all_bindings = editor:get_all_bindings()
    local matching = {}

    -- Normalize prefix (remove trailing dash if present)
    prefix = prefix:gsub("%-$", "")

    for _, binding in ipairs(all_bindings) do
        if starts_with(binding.sequence, prefix) and binding.sequence ~= prefix then
            local next_key = get_next_key(binding.sequence, prefix)
            if next_key then
                -- Check if we already have this next_key
                local found = false
                for _, m in ipairs(matching) do
                    if m.key == next_key then
                        found = true
                        break
                    end
                end

                if not found then
                    -- Determine if this is a prefix or a command
                    local is_prefix = false
                    local prefix_with_key = prefix .. " " .. next_key
                    for _, b in ipairs(all_bindings) do
                        if starts_with(b.sequence, prefix_with_key) and b.sequence ~= prefix_with_key then
                            is_prefix = true
                            break
                        end
                    end

                    table.insert(matching, {
                        key = next_key,
                        command = binding.command,
                        description = binding.description,
                        full_sequence = binding.sequence,
                        is_prefix = is_prefix,
                    })
                end
            end
        end
    end

    -- Sort by key
    table.sort(matching, function(a, b)
        return a.key < b.key
    end)

    return matching
end

-- Format hints into lines for display
function which_key.format_hints(bindings, max_width)
    max_width = max_width or 80
    local lines = {}
    local current_line = ""

    for i, binding in ipairs(bindings) do
        if i > which_key.config.max_hints then
            if current_line ~= "" then
                table.insert(lines, current_line)
            end
            table.insert(lines, "... and more")
            break
        end

        local hint
        if binding.is_prefix then
            hint = binding.key .. ":+"
        elseif which_key.config.show_description and binding.description ~= "" then
            -- Truncate long descriptions
            local desc = binding.description
            if #desc > 15 then
                desc = desc:sub(1, 12) .. "..."
            end
            hint = binding.key .. ":" .. desc
        else
            -- Truncate long command names
            local cmd = binding.command or "?"
            if #cmd > 15 then
                cmd = cmd:sub(1, 12) .. "..."
            end
            hint = binding.key .. ":" .. cmd
        end

        if current_line == "" then
            current_line = hint
        elseif #current_line + #which_key.config.separator + #hint <= max_width then
            current_line = current_line .. which_key.config.separator .. hint
        else
            table.insert(lines, current_line)
            current_line = hint

            if #lines >= which_key.config.max_lines then
                table.insert(lines, "...")
                return lines
            end
        end
    end

    if current_line ~= "" then
        table.insert(lines, current_line)
    end

    return lines
end

-- Show which-key hints in the extended echo area
function which_key.show_hints(prefix)
    if not which_key.config.enabled then
        return
    end

    local bindings = which_key.get_bindings_for_prefix(prefix)

    if #bindings > 0 then
        local lines = which_key.format_hints(bindings)
        -- Prepend the current prefix
        table.insert(lines, 1, prefix .. "-")
        editor:set_echo_area(lines)
        which_key.showing_hints = true
    end
end

-- Hide which-key hints
function which_key.hide_hints()
    if which_key.showing_hints then
        editor:clear_echo_area()
        which_key.showing_hints = false
    end
end

-- Check if we should show hints (called during idle)
function which_key.check_idle()
    if not which_key.config.enabled then
        return
    end

    -- Check if we're building a key sequence
    if editor:is_building_sequence() then
        local current = editor:current_sequence()
        -- Remove trailing dash for comparison
        current = current:gsub("%-$", "")

        -- Check if we've been idle long enough
        local idle_ms = editor:idle_time_ms()

        if idle_ms >= which_key.config.delay_ms then
            -- Show hints if we haven't already or prefix changed
            if not which_key.showing_hints or which_key.last_prefix ~= current then
                which_key.show_hints(current)
                which_key.last_prefix = current
            end
        end
    else
        -- Not building a sequence, hide hints if showing
        which_key.hide_hints()
        which_key.last_prefix = nil
    end
end

-- Register commands
editor:register_command("which-key", "Show available keybindings for prefix", function(args)
    if #args > 0 then
        which_key.show_hints(args[1])
    elseif editor:is_building_sequence() then
        local current = editor:current_sequence():gsub("%-$", "")
        which_key.show_hints(current)
    else
        -- Show all prefixes
        local all_bindings = editor:get_all_bindings()
        local prefixes = {}

        for _, binding in ipairs(all_bindings) do
            local parts = parse_sequence(binding.sequence)
            if #parts > 1 then
                local prefix = parts[1]
                if not prefixes[prefix] then
                    prefixes[prefix] = 0
                end
                prefixes[prefix] = prefixes[prefix] + 1
            end
        end

        local hints = {}
        for prefix, count in pairs(prefixes) do
            table.insert(hints, prefix .. ":+" .. count)
        end
        table.sort(hints)

        editor:message("Prefix keys: " .. table.concat(hints, " | "))
    end
    return {success = true}
end, {"wk"}, true, "s")

-- Describe key binding
editor:register_command("describe-key", "Describe a key binding", function(args)
    if #args == 0 then
        return {success = false, message = "Usage: describe-key <key-sequence>"}
    end

    local key = args[1]
    local all_bindings = editor:get_all_bindings()

    for _, binding in ipairs(all_bindings) do
        if binding.sequence == key then
            local msg = key .. " runs " .. binding.command
            if binding.description ~= "" then
                msg = msg .. " (" .. binding.description .. ")"
            end
            return {success = true, message = msg}
        end
    end

    -- Check if it's a prefix
    if editor:has_prefix_bindings(key) then
        which_key.show_hints(key)
        return {success = true}
    end

    return {success = false, message = key .. " is not bound"}
end, {"dk"}, true, "s")

-- Toggle which-key
editor:register_command("which-key-mode", "Toggle which-key hints", function(args)
    which_key.config.enabled = not which_key.config.enabled
    if which_key.config.enabled then
        editor:message("which-key enabled")
    else
        which_key.hide_hints()
        editor:message("which-key disabled")
    end
    return {success = true}
end)

-- Set delay
editor:register_command("which-key-delay", "Set which-key delay in milliseconds", function(args)
    if #args == 0 then
        return {success = true, message = "which-key delay: " .. which_key.config.delay_ms .. "ms"}
    end

    local delay = tonumber(args[1])
    if delay and delay >= 0 then
        which_key.config.delay_ms = delay
        return {success = true, message = "which-key delay set to " .. delay .. "ms"}
    else
        return {success = false, message = "Invalid delay value"}
    end
end, {}, true, "s")

-- Bind C-h k to describe-key (Emacs standard)
editor:bind_key("C-h k", "describe-key", "Describe a key binding")

-- Register the idle check function globally so the editor can call it
-- This gets called during the UI's idle loop
lumacs.which_key_check_idle = which_key.check_idle

-- Store module
lumacs.which_key = which_key

print("[which-key] Package loaded")

return which_key