cp.nvim/lua/cp/window.lua

---@class WindowState
---@field windows table<integer, WindowData>
---@field current_win integer
---@field layout string

---@class WindowData
---@field bufnr integer
---@field view table
---@field width integer
---@field height integer

local M = {}
local constants = require('cp.constants')

function M.clearcol()
    vim.api.nvim_set_option_value('number', false, { scope = 'local' })
    vim.api.nvim_set_option_value('relativenumber', false, { scope = 'local' })
    vim.api.nvim_set_option_value('statuscolumn', '', { scope = 'local' })
    vim.api.nvim_set_option_value('signcolumn', 'no', { scope = 'local' })
    vim.api.nvim_set_option_value('foldcolumn', '0', { scope = 'local' })
end

---@return WindowState
function M.save_layout()
    local windows = {}
    for _, win in ipairs(vim.api.nvim_list_wins()) do
        if vim.api.nvim_win_is_valid(win) then
            local bufnr = vim.api.nvim_win_get_buf(win)
            windows[win] = {
                bufnr = bufnr,
                view = vim.fn.winsaveview(),
                width = vim.api.nvim_win_get_width(win),
                height = vim.api.nvim_win_get_height(win),
            }
        end
    end

    return {
        windows = windows,
        current_win = vim.api.nvim_get_current_win(),
        layout = vim.fn.winrestcmd(),
    }
end

---@param state? WindowState
---@param tile_fn? fun(source_buf: integer, input_buf: integer, output_buf: integer)
function M.restore_layout(state, tile_fn)
    vim.validate({
        state = { state, { 'table', 'nil' }, true },
        tile_fn = { tile_fn, { 'function', 'nil' }, true },
    })

    if not state then
        return
    end

    vim.cmd.diffoff()

    local problem_id = vim.fn.expand('%:t:r')
    if problem_id == '' then
        for win, win_state in pairs(state.windows) do
            if
                vim.api.nvim_win_is_valid(win)
                and vim.api.nvim_buf_is_valid(win_state.bufnr)
            then
                local bufname = vim.api.nvim_buf_get_name(win_state.bufnr)
                if
                    not bufname:match('%.in$')
                    and not bufname:match('%.out$')
                    and not bufname:match('%.expected$')
                then
                    problem_id = vim.fn.fnamemodify(bufname, ':t:r')
                    break
                end
            end
        end
    end

    if problem_id ~= '' then
        vim.cmd('silent only')

        local base_fp = vim.fn.getcwd()
        local input_file = ('%s/io/%s.in'):format(base_fp, problem_id)
        local output_file = ('%s/io/%s.out'):format(base_fp, problem_id)
        local source_files = vim.fn.glob(problem_id .. '.*')
        local source_file
        if source_files ~= '' then
            local files = vim.split(source_files, '\n')
            local valid_extensions =
                vim.tbl_keys(constants.filetype_to_language)
            for _, file in ipairs(files) do
                local ext = vim.fn.fnamemodify(file, ':e')
                if vim.tbl_contains(valid_extensions, ext) then
                    source_file = file
                    break
                end
            end
            source_file = source_file or files[1]
        end

        if not source_file or vim.fn.filereadable(source_file) == 0 then
            return
        end

        vim.cmd.edit(source_file)
        local source_buf = vim.api.nvim_get_current_buf()
        local input_buf = vim.fn.bufnr(input_file, true)
        local output_buf = vim.fn.bufnr(output_file, true)

        if tile_fn then
            tile_fn(source_buf, input_buf, output_buf)
        else
            M.default_tile(source_buf, input_buf, output_buf)
        end
    else
        vim.cmd(state.layout)

        for win, win_state in pairs(state.windows) do
            if vim.api.nvim_win_is_valid(win) then
                vim.api.nvim_set_current_win(win)
                if vim.api.nvim_get_current_buf() == win_state.bufnr then
                    vim.fn.winrestview(win_state.view)
                end
            end
        end

        if vim.api.nvim_win_is_valid(state.current_win) then
            vim.api.nvim_set_current_win(state.current_win)
        end
    end
end

---@param source_buf integer
---@param input_buf integer
---@param output_buf integer
local function default_tile(source_buf, input_buf, output_buf)
    vim.validate({
        source_buf = { source_buf, 'number' },
        input_buf = { input_buf, 'number' },
        output_buf = { output_buf, 'number' },
    })

    vim.api.nvim_set_current_buf(source_buf)
    vim.cmd.vsplit()
    vim.api.nvim_set_current_buf(output_buf)
    vim.bo.filetype = 'cp'
    M.clearcol()
    vim.cmd(('vertical resize %d'):format(math.floor(vim.o.columns * 0.3)))
    vim.cmd.split()
    vim.api.nvim_set_current_buf(input_buf)
    vim.bo.filetype = 'cp'
    M.clearcol()
    vim.cmd.wincmd('h')
end

M.default_tile = default_tile

return M