table.lua

local Base = require('render-markdown.render.base')
local iter = require('render-markdown.lib.iter')
local log = require('render-markdown.core.log')
local str = require('render-markdown.lib.str')

---@class render.md.table.Data
---@field delim render.md.table.DelimRow
---@field rows render.md.table.Row[]

---@class render.md.table.DelimRow
---@field node render.md.Node
---@field cols render.md.table.DelimCol[]
---@field has_leading boolean
---@field has_trailing boolean

---@class render.md.table.DelimCol
---@field width integer
---@field alignment render.md.table.Alignment

---@enum render.md.table.Alignment
local Alignment = {
    left = 'left',
    right = 'right',
    center = 'center',
    default = 'default',
}

---@class render.md.table.Row
---@field node render.md.Node
---@field pipes render.md.Node[]
---@field cols render.md.table.Col[]
---@field has_leading boolean
---@field has_trailing boolean

---@class render.md.table.Col
---@field row integer
---@field start_col integer
---@field end_col integer
---@field width integer
---@field space render.md.table.Space

---@class render.md.table.Space
---@field left integer
---@field right integer

---@class render.md.render.Table: render.md.Render
---@field private config render.md.table.Config
---@field private data render.md.table.Data
local Render = setmetatable({}, Base)
Render.__index = Render

---@protected
---@return boolean
function Render:setup()
    self.config = self.context.config.pipe_table
    if not self.config.enabled then
        return false
    end

    -- ensure delimiter and rows exist
    local delim_node = nil ---@type render.md.Node?
    local row_nodes = {} ---@type render.md.Node[]
    local types = {
        delim = 'pipe_table_delimiter_row',
        row = { 'pipe_table_header', 'pipe_table_row' },
        skip = { 'block_continuation' },
    }
    self.node:for_each_child(function(node)
        if node.type == types.delim then
            delim_node = node
        elseif self.context.view:overlaps(node:get()) then
            if vim.tbl_contains(types.row, node.type) then
                row_nodes[#row_nodes + 1] = node
            elseif not vim.tbl_contains(types.skip, node.type) then
                log.unhandled(self.context.buf, 'markdown', 'row', node.type)
            end
        end
    end)
    if not delim_node or #row_nodes == 0 then
        return false
    end

    -- double check delimiter exists after parsing
    local delim = self:parse_delim(delim_node)
    if not delim then
        return false
    end

    -- double check rows exist after parsing
    local rows = {} ---@type render.md.table.Row[]
    table.sort(row_nodes)
    for _, row_node in ipairs(row_nodes) do
        local row = self:parse_row(row_node, #delim.cols)
        if row then
            rows[#rows + 1] = row
        end
    end
    if #rows == 0 then
        return false
    end

    -- store the max width in the delimiter
    for _, row in ipairs(rows) do
        for i, col in ipairs(row.cols) do
            local space = col.space.left + col.space.right
            local available = space - (2 * self.config.padding)
            -- if we don't have enough space for padding add it to the width
            local width = col.width
            if available < 0 then
                width = width - available
            end
            if self.config.cell == 'trimmed' then
                width = width - math.max(available, 0)
            end
            local delim_col = delim.cols[i]
            delim_col.width = math.max(delim_col.width, width)
        end
    end

    self.data = { delim = delim, rows = rows }

    return true
end

---@private
---@param node render.md.Node
---@return render.md.table.DelimRow?
function Render:parse_delim(node)
    local pipes, cells, has_leading, has_trailing =
        self:parse_cells(node, 'pipe_table_delimiter_cell')
    if not pipes or not cells then
        return nil
    end
    local cols = {} ---@type render.md.table.DelimCol[]

    ---@param start_col integer
    ---@param end_col integer
    ---@return integer
    local function get_width(start_col, end_col)
        local width = end_col - start_col
        assert(width >= 0, 'invalid table layout')
        if self.config.cell == 'padded' then
            width = math.max(width, self.config.min_width)
        elseif self.config.cell == 'trimmed' then
            width = self.config.min_width
        end
        return width
    end
    if not has_leading then
        -- as there is no leading pipe the first cell is used for alignment
        local start_col, end_col = cells[1].start_col, pipes[1].start_col
        -- add the first column
        cols[#cols + 1] = {
            width = get_width(start_col, end_col),
            alignment = Render.alignment(cells[1]),
        }
        -- now remove the first cell so that it isn't repocessed below
        table.remove(cells, 1)
    end
    for i, cell in ipairs(cells) do
        local start_col = pipes[i].end_col
        local end_col = 0
        if i + 1 <= #pipes then
            end_col = pipes[i + 1].start_col
        elseif not has_trailing then
            end_col = cells[i].end_col
        end
        cols[#cols + 1] = {
            width = get_width(start_col, end_col),
            alignment = Render.alignment(cell),
        }
    end
    ---@type render.md.table.DelimRow
    return {
        node = node,
        cols = cols,
        has_leading = has_leading,
        has_trailing = has_trailing,
    }
end

---@private
---@param node render.md.Node
---@return render.md.table.Alignment
function Render.alignment(node)
    local left = node:child('pipe_table_align_left')
    local right = node:child('pipe_table_align_right')
    if left and right then
        return Alignment.center
    elseif left then
        return Alignment.left
    elseif right then
        return Alignment.right
    else
        return Alignment.default
    end
end

---@private
---@param node render.md.Node
---@param num_cols integer
---@return render.md.table.Row?
function Render:parse_row(node, num_cols)
    local pipes, cells, has_leading, has_trailing =
        self:parse_cells(node, 'pipe_table_cell')
    if not pipes or not cells or #cells ~= num_cols then
        return nil
    end
    local cols = {} ---@type render.md.table.Col[]
    ---@param start_col integer
    ---@param end_col integer
    ---@param cell render.md.Node
    ---@return integer
    local function get_width(start_col, end_col, cell)
        local width = (end_col - start_col)
            - (cell.end_col - cell.start_col)
            + self.context:width(cell)
            + self.config.cell_offset({ node = cell:get() })
        assert(width >= 0, 'invalid table layout')
        return width
    end
    if not has_leading then
        local cell = cells[1]
        local start_col, end_col = cell.start_col, pipes[1].start_col
        cols[#cols + 1] = {
            row = cell.start_row,
            start_col = cell.start_col,
            end_col = cell.end_col,
            width = get_width(start_col, end_col, cell),
            space = {
                -- gap between the cell start and the pipe start
                left = math.max(cell.start_col - start_col, 0),
                -- attached to the end of the cell itself
                right = math.max(str.spaces('end', cell.text), 0),
            },
        }
        -- remove the first cell so that it is not reprocesed below
        table.remove(cells, 1)
    end
    for i, cell in ipairs(cells) do
        -- account for double width glyphs by replacing cell range with width
        local start_col = pipes[i].end_col
        local end_col = 0
        if i + 1 <= #pipes then
            end_col = pipes[i + 1].start_col
        elseif not has_trailing then
            end_col = cells[i].end_col
        end

        cols[#cols + 1] = {
            row = cell.start_row,
            start_col = cell.start_col,
            end_col = cell.end_col,
            width = get_width(start_col, end_col, cell),
            space = {
                -- gap between the cell start and the pipe start
                left = math.max(cell.start_col - start_col, 0),
                -- attached to the end of the cell itself
                right = math.max(str.spaces('end', cell.text), 0),
            },
        }
    end
    ---@type render.md.table.Row
    return {
        node = node,
        pipes = pipes,
        cols = cols,
        has_leading = has_leading,
        has_trailing = has_trailing,
    }
end

---@private
---@param node render.md.Node
---@param cell string
---@return render.md.Node[]?, render.md.Node[]?, boolean, boolean
function Render:parse_cells(node, cell)
    local pipes = {} ---@type render.md.Node[]
    local cells = {} ---@type render.md.Node[]
    local has_leading = true
    node:for_each_child(function(child)
        if child.type == '|' then
            pipes[#pipes + 1] = child
        elseif child.type == cell then
            cells[#cells + 1] = child
        else
            log.unhandled(self.context.buf, 'markdown', 'cell', child.type)
        end
        if #pipes == 0 and #cells == 1 then
            -- this means that the first cell was reached before the first cell (i.e. no leading pipe)
            has_leading = false
        end
    end)
    local has_trailing = pipes[#pipes].end_col >= cells[#cells].end_col
    local offset = 1 - (has_leading and 0 or 1) - (has_trailing and 0 or 1)
    if #pipes == 0 or #cells == 0 or #pipes ~= #cells + offset then
        return nil, nil, false, false
    end
    table.sort(pipes)
    table.sort(cells)
    return pipes, cells, has_leading, has_trailing
end

---@protected
function Render:run()
    self:delimiter()
    for _, row in ipairs(self.data.rows) do
        self:row(row)
    end
    if self.config.border_enabled then
        self:border()
    end
end

---@private
function Render:delimiter()
    local delim, border = self.data.delim, self.config.border

    local indicator, icon = self.config.alignment_indicator, border[11]
    local parts = iter.list.map(delim.cols, function(col)
        -- must have enough space to put the alignment indicator
        -- alignment indicator must be exactly one character wide
        -- do not put an indicator for default alignment
        local add_indicator = col.width >= 3
            and str.width(indicator) == 1
            and col.alignment ~= Alignment.default
        if not add_indicator then
            return icon:rep(col.width)
        end
        if col.alignment == Alignment.left then
            return indicator .. icon:rep(col.width - 1)
        elseif col.alignment == Alignment.right then
            return icon:rep(col.width - 1) .. indicator
        else
            return indicator .. icon:rep(col.width - 2) .. indicator
        end
    end)
    local delimiter = (delim.has_leading and border[4] or '')
        .. table.concat(parts, border[5])
        .. (delim.has_trailing and border[6] or '')

    local line = self:line()
    line:pad(str.spaces('start', delim.node.text))
    line:text(delimiter, self.config.head)
    line:pad(str.width(delim.node.text) - line:width())
    self.marks:over(self.config, 'table_border', delim.node, {
        virt_text = line:get(),
        virt_text_pos = 'overlay',
    })
end

---@private
---@param row render.md.table.Row
function Render:row(row)
    local icon = self.config.border[10]
    local header = row.node.type == 'pipe_table_header'
    local highlight = header and self.config.head or self.config.row

    if vim.tbl_contains({ 'trimmed', 'padded', 'raw' }, self.config.cell) then
        for _, pipe in ipairs(row.pipes) do
            self.marks:over(self.config, 'table_border', pipe, {
                virt_text = { { icon, highlight } },
                virt_text_pos = 'overlay',
            })
        end
    end

    if vim.tbl_contains({ 'trimmed', 'padded' }, self.config.cell) then
        for i, col in ipairs(row.cols) do
            local delim = self.data.delim.cols[i]
            local space = col.space
            local fill = delim.width - col.width
            -- delim(20) : --------------------
            -- col(4,7,2): ----XXXXXXX--
            -- fill(7)   :              _______
            if not self.context.conceal:enabled() then
                -- result: ----XXXXXXX--_______
                -- without concealing it is impossible to do full alignment
                self:shift(col, 'right', fill)
            elseif delim.alignment == Alignment.center then
                -- (7 + 2 - 4) // 2 = 5 // 2 = 2 -> move two spaces to the right
                -- result: __----XXXXXXX--_____
                local shift = math.floor((fill + space.right - space.left) / 2)
                self:shift(col, 'left', shift)
                self:shift(col, 'right', fill - shift)
            elseif delim.alignment == Alignment.right then
                -- 2 - 1 = 1 -> conceal one space on right side
                -- result: -_______----XXXXXXX-
                local shift = space.right - self.config.padding
                self:shift(col, 'left', fill + shift)
                self:shift(col, 'right', -shift)
            else
                -- 4 - 1 = 3 -> conceal three spaces on left side
                -- result: -XXXXXXX--_______---
                local shift = space.left - self.config.padding
                self:shift(col, 'left', -shift)
                self:shift(col, 'right', fill + shift)
            end
        end
    elseif self.config.cell == 'overlay' then
        self.marks:over(self.config, 'table_border', row.node, {
            virt_text = { { row.node.text:gsub('|', icon), highlight } },
            virt_text_pos = 'overlay',
        })
    end
end

---Use low priority to include pipe marks
---@private
---@param col render.md.table.Col
---@param side 'left'|'right'
---@param amount integer
function Render:shift(col, side, amount)
    local column = side == 'left' and col.start_col or col.end_col
    if amount > 0 then
        self.marks:add(self.config, true, col.row, column, {
            priority = 0,
            virt_text = self:line():pad(amount, self.config.filler):get(),
            virt_text_pos = 'inline',
        })
    elseif amount < 0 then
        amount = amount - self.context.conceal:width('', 1)
        self.marks:add(self.config, true, col.row, column + amount, {
            priority = 0,
            end_col = column,
            conceal = '',
        })
    end
end

---@private
function Render:border()
    local delim = self.data.delim
    local rows = self.data.rows
    local border = self.config.border

    ---@param row render.md.table.Row
    ---@return boolean
    local function width_equal(row)
        if vim.tbl_contains({ 'trimmed', 'padded' }, self.config.cell) then
            -- assume table was modified to match
            return true
        elseif self.config.cell == 'raw' then
            -- want the computed widths to match
            for i, col in ipairs(row.cols) do
                if delim.cols[i].width ~= col.width then
                    return false
                end
            end
            return true
        elseif self.config.cell == 'overlay' then
            -- want the underlying text widths to match
            return str.width(delim.node.text) == str.width(row.node.text)
        else
            return false
        end
    end

    local first, last = rows[1], rows[#rows]
    if not width_equal(first) or not width_equal(last) then
        return
    end

    ---@param node render.md.Node
    ---@return integer
    local function get_spaces(node)
        local _, line = node:line('first', 0)
        return math.max(str.spaces('start', line or ''), node.start_col)
    end

    local first_node = first.node
    local last_node = #rows == 1 and delim.node or last.node
    local spaces = get_spaces(first_node)
    if spaces ~= get_spaces(last_node) then
        return
    end

    local sections = iter.list.map(delim.cols, function(col)
        return border[11]:rep(col.width)
    end)

    ---@param node render.md.Node
    ---@param above boolean
    ---@param chars { [1]: string, [2]: string, [3]: string }
    local function table_border(node, above, chars)
        local text = chars[1] .. table.concat(sections, chars[2]) .. chars[3]
        local highlight = above and self.config.head or self.config.row
        local line = self:line():pad(spaces):text(text, highlight)

        local virtual = self.config.border_virtual
        local row, target = node:line(above and 'above' or 'below', 1)
        local available = target and str.width(target) == 0

        if not virtual and available and self.context.used:take(row) then
            self.marks:add(self.config, 'table_border', row, 0, {
                virt_text = line:get(),
                virt_text_pos = 'overlay',
            })
        else
            self.marks:add(self.config, 'virtual_lines', node.start_row, 0, {
                virt_lines = { self:indent():line(true):extend(line):get() },
                virt_lines_above = above,
            })
        end
    end
    table_border(first_node, true, {
        first.has_leading and border[1] or '',
        border[2],
        first.has_trailing and border[3] or '',
    })
    if #rows > 1 then
        table_border(last_node, false, {
            first.has_leading and border[7] or '',
            border[8],
            first.has_trailing and border[9] or '',
        })
    end
end

return Render