Module:Table

---@module 'Table'
--- Parse and manipulate wikitext tables in MediaWiki.
---
--- Provides functions to extract tables, parse them into structured data,
--- and build a slot grid accounting for colspan/rowspan, classes, and styles.

local _M = {}

local _gsub = mw.ustring.gsub
local _sub = mw.ustring.sub
local _match = mw.ustring.match
local _len = mw.ustring.len
local _gmatch = mw.ustring.gmatch
local _gsplit = mw.text.gsplit
local _tostring = tostring
local table_insert = table.insert
local _tonumber = tonumber

--- Internal cache for strings that have already been trimmed
local trim_cache = {}
--- Internal reference of what strings are to be considered whitespace for wikitext conversion
local whitespace = {
    [' '] = true,
    ['\n'] = true,
    ['\t'] = true,
    ['\r'] = true,
}

---
--- Error logging
---
---@param msg string Error message to log
---@param where 'console'|'preview' Where to log the error: 'console' or 'preview'
---@return nil
local function add_error(msg, where)
    if where == 'console' then
        mw.log('Module:Table error: ' .. msg)
    elseif where == 'preview' then
        mw.addWarning('<span class="error"><strong>[[Module:Table]] error:</strong>&nbsp;' .. msg .. '</span>')
    end
    return nil
end

---
--- Protected call utility
---
---@param fn function Function to call
---@param ... any Arguments to pass to the function
---@return any|nil output Result of the function call, or `nil` if an error occurred
local function try_call(fn, ...)
    local ok, output = xpcall(fn, function(err)
        add_error('Unexpected error in <code>try_call()</code>: ' .. _tostring(err), 'console')
    end, ...)
    if ok then
        return output
    else
        return nil
    end
end

---
---Convert to integer >= 0
---
---@param input any Input to convert
---@return integer|nil integer Non-negative integer or `nil` if invalid
function _M.to_integer(input)
    local num = _tonumber(input)
    if num and num >= 0 and math.floor(num) == num then
        return num
    end
    add_error('Expected non-negative integer but got: ' .. _tostring(input), 'console')
    return nil
end

---
--- Finds first non-whitespace character in a string
---
---@param s string Input string
---@param len integer Length of the string
---@return integer|nil index Index of the first non-whitespace character, or `nil`
local function find_first_nonwhitespace(s, len)
    for i = 1, len do
        if not whitespace[_sub(s, i, i)] then
            return i
        end
    end
end

---
--- Finds last non-whitespace character in a string
---
---@param s string Input string
---@param len integer Length of the string
---@return integer|nil index Index of the last non-whitespace character, or `nil`
local function find_last_nonwhitespace(s, len)
    for i = len, 1, -1 do
        if not whitespace[_sub(s, i, i)] then
            return i
        end
    end
end

---
--- Trims leading and trailing whitespace from a string
---
---@param s any Input string
---@return string trimmed Trimmed string
function _M.trim_whitespace(s)
    local len = _len(s)
    local low = find_first_nonwhitespace(s, len)
    if not low then
        return ''
    end
    local high = find_last_nonwhitespace(s, len)
    if not high then
        add_error('Unexpected end in <code>trim_whitespace()</code> for input: ' .. _tostring(s), 'console')
        return ''
    end
    return _sub(s, low, high)
end

-- Cached fast trim
function _M.cheap_trim(input)
    if trim_cache[input] then
        return trim_cache[input]
    end
    local trimmed = _M.trim_whitespace(input)
    trim_cache[input] = trimmed
    return trimmed
end

---
--- Parse a single cell
---
function _M.parse_cell(cell_wikitext)
    local cell = {}
    cell.colspan = _tonumber(_match(cell_wikitext, 'colspan *= *"?([0-9]+)"?')) or 1
    cell.rowspan = _tonumber(_match(cell_wikitext, 'rowspan *= *"?([0-9]+)"?')) or 1
    cell.text = _gsub(cell_wikitext, 'colspan *= *"?[0-9]+"?', "")
    cell.text = _gsub(cell.text, 'rowspan *= *"?[0-9]+"?', "")
    cell.text = _M.cheap_trim(cell.text)
    return cell
end

-- Extract tables from wikitext safely
function _M.get_tables(wikitext)
    local tables = {}
    wikitext = '\n' .. wikitext
    for t in _gmatch(wikitext, '\n{|.-\n|}') do
        table_insert(tables, _M.cheap_trim(t))
    end
    return tables
end

---
---Get table by ID attribute
---
---@param wikitext string
---@param id string
---@return string|nil wikitext Wikitext with the specified ID, or `nil` if not found
function _M.get_table_by_id(wikitext, id)
    local value
    local tables = _M.get_tables(wikitext)
    for _, t in ipairs(tables) do
        local value = _match(t, "^{|[^\n]*id *= *[\"']?([^\"'\n]+)[\"']?[^\n]*\n")
        if not value then
            value = _match(t, "^{|[^\n]*id *= *[\'']?([^\''\n]+)[\'']?[^\n]*\n")
        end
        if value == id then
            return t
        end
    end
    return nil
end

---
---Parse table wikitext into structured data
---
---@param table_wikitext any
---@return table table_data Table data as a list of rows, each containing a list of cell objects
function _M.get_table_data(table_wikitext)
    local table_data = {}
    local text = _M.cheap_trim(table_wikitext)

    text = _gsub(text, "^{|.-\n", "")
    text = _gsub(text, "\n|}$", "")
    text = _gsub(text, "^|%+.-\n", "")
    text = _gsub(text, "|%-.-\n", "|-\n")
    text = _gsub(text, "^|%-\n", "")
    text = _gsub(text, "\n|%-$", "")

    for row_wikitext in _gsplit(text, '|-', true) do
        local row_data = {}
        row_wikitext = _gsub(row_wikitext, '||', '\n|')
        row_wikitext = _gsub(row_wikitext, '!!', '\n|')
        row_wikitext = _gsub(row_wikitext, '\n!', '\n|')
        row_wikitext = _gsub(row_wikitext, '^!', '\n|')
        row_wikitext = _gsub(row_wikitext, '^\n|', '')

        for cell_wikitext in _gsplit(row_wikitext, "\n|") do
            if cell_wikitext ~= '' then
                table_insert(row_data, _M.parse_cell(cell_wikitext))
            end
        end
        if #row_data > 0 then
            table_insert(table_data, row_data)
        end
    end
    return table_data
end

---
---Build slot grid
---Accounts for colspan and rowspan, fills in `nil` for empty slots.
---
---@param table_data table Table data as returned by `get_table_data()`
---@return table slots 2D array representing the slot grid with merged cells accounted for
function _M.get_table_slots(table_data)
    if not table_data or type(table_data) ~= 'table' then
        add_error('Invalid table: must be a table of rows', 'console')
        return {}
    end

    local slots = {}

    for rowIndex, row in ipairs(table_data) do
        if type(row) ~= 'table' then
            add_error('Invalid row at index ' .. rowIndex .. ': must be a table of cells', 'console')
        else
            for cellIndex, cell in ipairs(row) do
                if type(cell) ~= 'table' then
                    add_error('Invalid cell at row ' .. rowIndex .. ', column ' .. cellIndex, 'console')
                else
                    local rowspan = cell.rowspan or 1
                    local colspan = cell.colspan or 1
                    local x = cellIndex
                    local y = rowIndex

                    -- Skip occupied slots (from previous rowspan/colspan)
                    while slots[y] and slots[y][x] do
                        x = x + 1
                    end

                    -- Fill slots
                    for dy = 0, rowspan - 1 do
                        for dx = 0, colspan - 1 do
                            while (y + dy) > #slots do
                                table_insert(slots, {})
                            end
                            slots[y + dy][x + dx] = cell
                        end
                    end
                end
            end
        end
    end

    return slots
end

---
---Render slot grid into wikitable syntax
---Preserves merged cell logic, skips nil slots, only outputs each cell once,
---can add styles, and set `colspan` and `rowspan` per cell.
---
---@param slots table Slot grid as returned by `get_table_slots()`
---@param cell_class_function? function Optional function to generate additional attributes for each cell. It should accept three parameters: the cell `object`, its row index (`y_axis`), and its column index (`x_axis`). It should return a string of additional attributes (e.g., `'class="my-class" style="color: red;"'`) or an empty string `''` if no additional attributes are needed.
---@param table_class? string Optional class(es) attribute for the entire table (default: `wikitable`)
---@return string wikitext_table Wikitext table output
function _M.render_slots(slots, cell_class_function, table_class)
    table_class = table_class or 'wikitable'
    local output = { '{| class="' .. table_class .. '"' }
    local used = {}

    for y_axis, row in ipairs(slots) do
        table_insert(output, '|-')
        for x_axis, cell in ipairs(row) do
            if cell and not used[cell] then
                used[cell] = true
                local parts = {}

                if cell.rowspan and cell.rowspan > 1 then
                    table_insert(parts, 'rowspan = ' .. _tostring(cell.rowspan))
                end
                if cell.colspan and cell.colspan > 1 then
                    table_insert(parts, 'colspan = ' .. _tostring(cell.colspan))
                end

                if cell_class_function then
                    local custom_attr = cell_class_function(cell, y_axis, x_axis)
                    if custom_attr and custom_attr ~= '' then
                        table_insert(parts, custom_attr)
                    end
                end

                local attr_str = (#parts > 0) and (table.concat(parts, ' ') .. ' |') or '|'
                table_insert(output, attr_str .. ' ' .. (cell.text or ''))
            end
        end
    end

    table_insert(output, '|}')
    return table.concat(output, '\n')
end

---
--- Convenience: get slot grid by table ID
---
function _M.slots_from_wikitext_by_id(wikitext, id)
    local t = _M.get_table_by_id(wikitext, id)
    if not t then
        return nil
    end
    return _M.get_table_slots(_M.get_table_data(t))
end

---
--- Parses flat JSON object into structured table data
---
--- e.g. `[{"col1":"A","col2":"B"},{"col1":"C","col2":"D"}]`
--- into
--- ```lua
--- {
---   { {text="A", colspan=1, rowspan=1}, {text="B", colspan=1, rowspan=1} },
---   { {text="C", colspan=1, rowspan=1}, {text="D", colspan=1, rowspan=1} }
--- }
--- ```
---@param data table Flat JSON object as parsed by `parse_json()`
---@return table|nil table_data Structured table data or `nil` if invalid
local function flat_json_to_table(data)
    if type(data) ~= 'table' then
        add_error('Invalid data: must be a list of tables', 'preview')
        return nil
    end
    -- If the first element is a table, assume it's already in the structured format
    if #data > 0 and type(data[1]) == 'table' then
        return data
    end
    local table_data = {}
    for _, row in ipairs(data) do
        if type(row) == 'table' then
            local row_data = {}
            -- Note: pairs() will iterate over string keys (like 'col1', 'col2')
            for _, cell in pairs(row) do
                table_insert(row_data, {
                    text = _tostring(cell),
                    colspan = 1,
                    rowspan = 1
                })
            end
            table_insert(table_data, row_data)
        else
            add_error('Invalid row in flat JSON object removed. Must be a table: ' .. _tostring(row_data), 'preview')
            -- This break is problematic if only one row is bad, but keeping original logic flow
            break
        end
    end
    -- FIX: Should return the constructed table_data, not an undefined 'slots'
    return table_data
end

---
--- Parse JSON-style input safely
--- Will also parse if Lua table syntax is used instead of JSON
---
---@param json_data string JSON data to parse
---@return table|nil parsed_json_output Parsed JSON data or `nil` if invalid
function _M.parse_json(json_data)
    -- First attempt: Strict JSON (standard mw.text.jsonDecode)
    local ok_json, data = pcall(mw.text.jsonDecode, json_data)
    if ok_json and type(data) == 'table' then
        return data
    end

    -- Second attempt: Try to convert to loose Lua table/JSON syntax
    -- 1. Replace single quotes with double quotes (handles 'string' -> "string")
    local _json_data = _gsub(json_data, "'", '"')
    
    -- 2. Quote unquoted keys (handles {key: value} and {key = value} -> {"key": value})
    -- FIX: Use a character set [:=] to match both colon (JSON) and equals (Lua) delimiters
    _json_data = _gsub(_json_data, '([%w_]+)%s*([:=])', '"%1"%2')

    local ok_lua, data_lua = pcall(mw.text.jsonDecode, _json_data)
    if ok_lua and type(data_lua) == 'table' then
        return data_lua
    end

    -- If both fail, log the error with the original data
    add_error('Invalid input: must be JSON or Lua table. Original input failed to parse.', 'preview')
    return nil
end


    --local tbl_data = data:gsub("'", '"'):gsub("([%w_]+)%s*:", '"%1":')
    --if not ok or type(data) ~= 'table' then
    --    add_error(mw.ustring.format(
    --            'Invalid JSON input. Input: %s. Output: %s.',
    --            data,
    --            tbl_data),
    --        'preview')
    --    return nil
    --end
    --ok, tbl_data2 = pcall(mw.text.jsonDecode, tbl_data)
    --if type(tbl_data2) == 'table' and ok then
    --    return tbl_data2
    --else
    --    local ok2, tbl_data3 = pcall(mw.text.jsonDecode, json_data)
    --    if ok2 and type(tbl_data3) == 'table' then
    --        return tbl_data3
    --    end
    --end
    --add_error('Invalid input: must be JSON or Lua table.', 'preview')
    --return data
--end

---
--- Fetch and clean arguments passed as input by removing blank arguments,
--- trimming whitespace, and in wrapper templates only checking parentArgs
--- for efficiency.
---
---@param frame table `frame` object containing `args` field as received during invocation of the module with `{{#invoke:...}}` or from a template transclusion. Note, this will be a `table` object if called from another module.
---@return table<integer|string, any> clean_args Clean arguments after parsing with input in key-value pairs e.g. `{ param1 = 'value1', param2 = 'value2' }`
local function fetch_args(frame)
    local getArgs = require('Module:Arguments').getArgs
    if frame and frame.args then
        return getArgs(frame, {
            removeBlanks = true,
            trim = true,
            wrappers = { 'Template:Table', 'Template:Table/sandbox' }
        })
    else
        add_error('No argument(s) passed to module.', 'preview')
        return {}
    end
end

---
--- Build table from `frame.args` in JSON format
---
--- #### Templates
-- Templates should use the `build()` function as an entry point.<br/>```{{#invoke: Table | build | ... }}```
--- #### Modules
--- Based on input:
--- - input is already a fully-expanded slot grid, use `render_slots()` directly.
--- - input is structured data (rows/cells), run it through `get_table_slots()` first.
--- - is a flat JSON object like `[{"col1":"A","col2":"B"}]`, use `parse_json()`.
---
---@param frame table Frame object with an ['args'] field, or table, containing the input.
---@return string wikitext_table Wikitext table output
function _M.build(frame)
    local args = fetch_args(frame)
    local data = args.data and _M.parse_json(args.data)
    if not data then
        return add_error(
            mw.ustring.format(
                'No output after argument(s) parsed as JSON. See [[Module:Table|template documentation]]. Input: %s. Output: %s.',
                args.data or 'nil', data or 'nil'),
            'preview')
    end
    local table_data = flat_json_to_table(data) or data
    if not table_data then
        return add_error(
            mw.ustring.format(
                'No output after JSON output parsed into table data. See [[Module:Table|template documentation]]. Input: %s and %s. Output: %s.',
                args.data or 'nil', data or 'nil', table_data or 'nil'),
            'preview')
    end
    local slots = _M.get_table_slots(table_data)
    return _M.render_slots(slots)
end

return _M