-- lua filter for RST-like list-tables in Markdown.
-- Copyright (C) 2021 Martin Fischer, released under MIT license

if PANDOC_VERSION and PANDOC_VERSION.must_be_at_least then
    PANDOC_VERSION:must_be_at_least("2.11")
else
    error("pandoc version >=2.11 is required")
end

-- Get the list of cells in a row.
local row_cells = function (row) return row.cells or {} end

-- "Polyfill" for older pandoc versions.
if PANDOC_VERSION <= '2.16.2' then
  -- previous pandoc versions used simple Attr/list pairs
  pandoc.Row = function (cells) return {{}, cells} end
  pandoc.TableHead = function (rows) return {{}, rows or {}} end
  pandoc.TableFoot = function (rows) return {{}, rows or {}} end
  pandoc.Cell = function (contents, align, rowspan, colspan, attr)
    return {
      attr = attr or pandoc.Attr(),
      alignment = align or pandoc.AlignDefault,
      contents = contents or {},
      col_span = colspan or 1,
      row_span = rowspan or 1
    }
  end
  row_cells = function (row) return row[2] end
end

local alignments = {
    d = 'AlignDefault',
    l = 'AlignLeft',
    r = 'AlignRight',
    c = 'AlignCenter'
}


-- This is like assert() but it can take a Block or Blocks 'where' argument
-- and will output the corresponding markdown (truncated at 1024 characters).
local function assert_(assertion, message, where)
    message = message or 'assertion failed!'
    if not assertion then
        local extra = ''
        if where then
            local blocks = pandoc.Blocks(where)
            local markdown = pandoc.write(pandoc.Pandoc(blocks), 'markdown')
            extra = ' at\n' .. markdown:sub(1, 1024) ..
                (#markdown > 1024 and '...' or '')
        end
        error(message .. extra, 2)
    end
end

local function get_colspecs(div_attributes, column_count)
    -- list of (align, width) pairs
    local colspecs = {}

    for i = 1, column_count do
        table.insert(colspecs, {pandoc.AlignDefault, nil})
    end

    if div_attributes.aligns then
        local i = 1
        for a in div_attributes.aligns:gmatch('[^,]+') do
            assert_(alignments[a] ~= nil,
                    "unknown column alignment " .. tostring(a))
            colspecs[i][1] = alignments[a]
            i = i + 1
        end
        div_attributes.aligns = nil
    end

    if div_attributes.widths then
        local total = 0
        local widths = {}
        for w in div_attributes.widths:gmatch('[^,]+') do
            table.insert(widths, tonumber(w))
            total = total + tonumber(w)
        end
        for i = 1, column_count do
            colspecs[i][2] = widths[i] / total
        end
        div_attributes.widths = nil
    end

    return colspecs
end

local function new_table_body(rows, attr, header_col_count)
    attr = attr or {}
    return {
        attr = attr,
        body = rows,
        head = {},
        row_head_columns = header_col_count
    }
end

local function new_cell(contents)
    local attr = {}
    local colspan = 1
    local rowspan = 1
    local align = pandoc.AlignDefault

    -- At the time of writing this Pandoc does not support attributes
    -- on list items, so we use empty spans as a workaround.
    if contents[1] and contents[1].content then
        if contents[1].content[1] and contents[1].content[1].t == "Span" then
            if #contents[1].content[1].content == 0 then
                attr = contents[1].content[1].attr
                table.remove(contents[1].content, 1)
                colspan = attr.attributes.colspan or 1
                attr.attributes.colspan = nil
                rowspan = attr.attributes.rowspan or 1
                attr.attributes.rowspan = nil
                align = alignments[attr.attributes.align] or pandoc.AlignDefault
                attr.attributes.align = nil
            end
        end
    end

    return pandoc.Cell(contents, align, rowspan, colspan, attr)
end

local function process(div)
    if (div.attr.classes[1] ~= "list-table" and
        div.attr.classes[1] ~= "list-table-body") then return nil end
    local class = div.attr.classes[1]
    table.remove(div.attr.classes, 1)

    if #div.content == 0 then return nil end

    local caption = {}

    if div.content[1].t == "Para" then
        local para = table.remove(div.content, 1)
        caption = {pandoc.Plain(para.content)}
    end

    if #div.content == 0 then return nil end

    assert_(div.content[1].t == "BulletList",
            "expected bullet list, found " .. div.content[1].t, div.content[1])
    local list = div.content[1]

    -- rows points to the current body's rows
    local bodies = {attr=nil, {rows={}}}
    local rows = bodies[#bodies].rows

    for i = 1, #list.content do
        local attr = nil
        if (#list.content[i] > 1) then
            assert_(list.content[i][1].content, "expected list item to " ..
                        "have a content attr", list.content[i][1])
            assert_(#list.content[i][1].content == 1, "expected row attrs " ..
                        "to contain only one inline",
                    list.content[i][1].content)
            assert_(list.content[i][1].content[1].t == "Span", "expected " ..
                        "row attrs to contain a span",
                    list.content[i][1].content[1])
            assert_(#list.content[i][1].content[1].content == 0, "expected " ..
                        "row attrs span to be empty",
                    list.content[i][1].content[1])
            attr = list.content[i][1].content[1].attr
            table.remove(list.content[i], 1)
        end

        assert_(#list.content[i] == 1, "expected item to contain only one " ..
                    "block", list.content[i])

        if (list.content[i][1].t ~= 'Table') then
            assert_(list.content[i][1].t == "BulletList",
                    "expected bullet list, found " .. list.content[i][1].t,
                    list.content[i][1])
            local cells = {}
            for _, cell_content in pairs(list.content[i][1].content) do
                table.insert(cells, new_cell(cell_content))
            end
            local row = pandoc.Row(cells, attr)
            table.insert(rows, row)

        else
            local tab = list.content[i][1]
            -- XXX is there a better way to check that there's no caption?
            assert_((not tab.caption.long or #tab.caption.long == 0) and
                        (not tab.caption.short or #tab.caption.short == 0),
                    "table bodies can't have captions (they'd be " ..
                        "ignored)", tab)
            -- XXX would have to check against default colspecs to know whether
            --     any have been defined?
            -- assert_(#tab.colspecs == 0, "table bodies can't (yet) have " ..
            --         "column specs", tab)
            assert_(#tab.head.rows == 0, "table bodies can't (yet) have " ..
                        "headers", tab)
            assert_(#tab.bodies == 1, "table bodies can't contain other " ..
                        "table bodies", tab)

            if #rows > 0 then
                table.insert(bodies, {attr=nil, rows={}})
                rows = bodies[#bodies].rows
            end

            bodies[#bodies].attr = tab.attr
            for _, row in ipairs(tab.bodies[1].body) do
                table.insert(rows, row)
            end
        end
    end

    -- switch back to the first body
    rows = bodies[1].rows

    local header_row_count = tonumber(div.attr.attributes['header-rows']) or
        (class == 'list-table' and 1 or 0)
    div.attr.attributes['header-rows'] = nil

    local header_col_count = tonumber(div.attr.attributes['header-cols']) or 0
    div.attr.attributes['header-cols'] = nil

    local column_count = 0
    for i = 1, #row_cells(rows[1] or {}) do
        column_count = column_count + row_cells(rows[1])[i].col_span
    end

    local colspecs = get_colspecs(div.attr.attributes, column_count)
    local thead_rows = {}
    for i = 1, header_row_count do
        table.insert(thead_rows, table.remove(rows, 1))
    end

    local new_bodies = {}
    for _, body in ipairs(bodies) do
        if #body.rows > 0 then
            table.insert(new_bodies, new_table_body(body.rows, body.attr,
                                                    header_col_count))
        end
        -- XXX this should be a body property
        header_col_count = 0
    end

    return pandoc.Table(
        {long = caption, short = {}},
        colspecs,
        pandoc.TableHead(thead_rows),
        new_bodies,
        pandoc.TableFoot(),
        div.attr
    )
end

return {{Div = process}}
