-- latex-tables.lua — write tables in LaTeX, render them in HTML too.
--
--   \begin{tabular}{|l|c|r|}
--   \hline
--   Quantity & Symbol & Unit \\
--   \hline
--   pressure & $p$ & Pa \\
--   density  & $\rho$ & kg/m$^3$ \\
--   \hline
--   \end{tabular}
--
-- Supported: l c r, p{w} m{w} b{w}, X (tabularx), *{n}{...} repetition,
-- @{...} and >{...}<{...} (ignored), | vertical rules, \hline, \cline{a-b},
-- booktabs \toprule \midrule \bottomrule \cmidrule{a-b}, \multicolumn,
-- \multirow, and the \begin{table} float wrapper with \caption and \label.
--
-- PDF: passed through untouched (it is already LaTeX).
-- HTML: converted to a real <table>; cell contents are parsed as LaTeX, so
--       $math$, \textbf{...} etc. work.
--
-- Header rows: everything above the first rule that follows row 1 becomes
-- <thead>. Put "% no-header" inside the environment to disable that.

local ENVS = { tabular = true, tabularx = true, longtable = true,
               ["tabular*"] = true }

local function fmt_is(name)
  if quarto and quarto.doc and quarto.doc.is_format then
    return quarto.doc.is_format(name)
  end
  return FORMAT ~= nil and FORMAT:match(name) ~= nil
end

------------------------------------------------------------------ utils

local function trim(s)
  return (s:gsub("^%s+", ""):gsub("%s+$", ""))
end

-- index of the "}" matching the "{" at open_idx
local function find_close(s, open_idx)
  local depth, i, n = 0, open_idx, #s
  while i <= n do
    local c = s:sub(i, i)
    if c == "\\" then
      i = i + 2
    else
      if c == "{" then
        depth = depth + 1
      elseif c == "}" then
        depth = depth - 1
        if depth == 0 then return i end
      end
      i = i + 1
    end
  end
  return nil
end

-- split on every literal `sep` sitting at brace depth 0 and outside any
-- nested environment (so \\ and & inside \begin{cases}/{array}/{pmatrix}
-- in a cell are left alone)
local function split_top(s, sep)
  local parts, depth, env, i, start = {}, 0, 0, 1, 1
  local n, k = #s, #sep
  while i <= n do
    if depth == 0 and env == 0 and s:sub(i, i + k - 1) == sep then
      parts[#parts + 1] = s:sub(start, i - 1)
      i = i + k
      start = i
    else
      local c = s:sub(i, i)
      if c == "\\" then
        if s:find("^\\begin%s*{", i) then
          env = env + 1
          i = i + 6
        elseif s:find("^\\end%s*{", i) then
          env = math.max(0, env - 1)
          i = i + 4
        else
          i = i + 2
        end
      else
        if c == "{" then
          depth = depth + 1
        elseif c == "}" then
          depth = math.max(0, depth - 1)
        end
        i = i + 1
      end
    end
  end
  parts[#parts + 1] = s:sub(start)
  return parts
end

-- strip % comments (but keep \%)
local function strip_comments(s)
  local out, i, n = {}, 1, #s
  while i <= n do
    local c = s:sub(i, i)
    if c == "\\" then
      out[#out + 1] = s:sub(i, i + 1)
      i = i + 2
    elseif c == "%" then
      local nl = s:find("\n", i, true)
      if nl == nil then break end
      i = nl -- keep the newline
    else
      out[#out + 1] = c
      i = i + 1
    end
  end
  return table.concat(out)
end

-- LaTeX length -> CSS length
local function css_len(w)
  w = trim(w)
  local frac, rel = w:match("^([%d%.]+)\\(%a+)$")
  if frac ~= nil and (rel == "textwidth" or rel == "linewidth"
      or rel == "columnwidth") then
    return string.format("%.4g%%", tonumber(frac) * 100)
  end
  if w:match("^\\%a+$") then return "100%" end
  local num, unit = w:match("^([%-%d%.]+)%s*(%a+)$")
  if num ~= nil then
    if unit == "pt" then return num .. "pt" end
    if unit == "cm" or unit == "mm" or unit == "in"
        or unit == "em" or unit == "ex" then
      return num .. unit
    end
  end
  return nil
end

------------------------------------------------------------------ colspec

-- *{3}{c|} -> c|c|c|
local function expand_star(spec)
  local guard = 0
  while guard < 20 do
    guard = guard + 1
    local st, _, cnt, body = spec:find("%*%s*(%b{})%s*(%b{})")
    if st == nil then break end
    local n = tonumber(cnt:sub(2, -2))
    local inner = body:sub(2, -2)
    if n == nil then break end
    spec = spec:sub(1, st - 1) .. string.rep(inner, n)
        .. spec:sub(st + #cnt + #body + (spec:sub(st + 1, st + 1) == " " and 1 or 0))
    -- recompute safely: rebuild from the matched span
    local _, en = spec:find(inner, st, true)
    if en == nil then break end
  end
  return spec
end

-- returns list of { align, left, right, width }
local function parse_colspec(spec)
  spec = expand_star(spec)
  local cols, pending, i, n = {}, 0, 1, #spec
  while i <= n do
    local c = spec:sub(i, i)
    if c == "|" then
      pending = pending + 1
      i = i + 1
    elseif c == " " or c == "\n" or c == "\t" then
      i = i + 1
    elseif c == "@" or c == "!" or c == ">" or c == "<" then
      local br = spec:find("{", i, true)
      local close = br and find_close(spec, br)
      i = close and (close + 1) or (i + 1)
    elseif c == "l" or c == "c" or c == "r" then
      cols[#cols + 1] = { align = c, left = pending, right = 0 }
      pending = 0
      i = i + 1
    elseif c == "p" or c == "m" or c == "b" then
      local br = spec:find("{", i, true)
      local close = br and find_close(spec, br)
      local w = close and spec:sub(br + 1, close - 1) or nil
      cols[#cols + 1] = { align = "l", left = pending, right = 0,
                          width = w and css_len(w) or nil }
      pending = 0
      i = close and (close + 1) or (i + 1)
    elseif c == "X" then
      cols[#cols + 1] = { align = "l", left = pending, right = 0, width = "auto" }
      pending = 0
      i = i + 1
    else
      i = i + 1
    end
  end
  if #cols > 0 and pending > 0 then cols[#cols].right = pending end
  return cols
end

------------------------------------------------------------------ rows

local RULE_PAT = {
  { pat = "^\\toprule%s*(%b[])?", kind = "thick" },
  { pat = "^\\toprule", kind = "thick" },
  { pat = "^\\bottomrule", kind = "thick" },
  { pat = "^\\midrule", kind = "thin" },
  { pat = "^\\hline", kind = "thin" },
}

-- pull leading rule commands off a row segment
local function take_rules(seg)
  local rules = { all = nil, ranges = {}, count = 0 }
  local changed = true
  while changed do
    changed = false
    seg = seg:gsub("^%s+", "")
    -- \noalign{...} / \addlinespace / spacing args: drop
    local br = seg:match("^\\noalign%s*%b{}")
    if br then seg = seg:sub(#br + 1); changed = true end
    local al = seg:match("^\\addlinespace%s*%b[]") or seg:match("^\\addlinespace")
    if al then seg = seg:sub(#al + 1); changed = true end
    local opt = seg:match("^%b[]")
    if opt then seg = seg:sub(#opt + 1); changed = true end
    for _, r in ipairs(RULE_PAT) do
      local m = seg:match(r.pat)
      if m then
        rules.count = rules.count + 1
        rules.all = (rules.all == "thick" or r.kind == "thick") and "thick" or "thin"
        if rules.count > 1 then rules.all = "double" end
        seg = seg:sub(#m + 1)
        changed = true
        break
      end
    end
    -- \cline{a-b} and \cmidrule(lr){a-b}
    local head, range = seg:match("^(\\cline%s*(%b{}))")
    if head == nil then
      local h2, _, r2 = seg:match("^(\\cmidrule%s*(%b())%s*(%b{}))")
      if h2 then head, range = h2, r2 end
    end
    if head == nil then
      local h3, r3 = seg:match("^(\\cmidrule%s*(%b{}))")
      if h3 then head, range = h3, r3 end
    end
    if head ~= nil and range ~= nil then
      local a, b = range:sub(2, -2):match("^%s*(%d+)%s*%-%s*(%d+)%s*$")
      if a then
        rules.ranges[#rules.ranges + 1] = { tonumber(a), tonumber(b) }
      end
      seg = seg:sub(#head + 1)
      changed = true
    end
  end
  return seg, rules
end

local function has_rule(rules)
  return rules.all ~= nil or #rules.ranges > 0
end

-- \multicolumn{n}{spec}{content} / \multirow{n}{w}{content}
local function parse_cell(txt)
  local cell = { colspan = 1, rowspan = 1, body = txt, spec = nil }
  local t = trim(txt)
  local head, a, b = t:match("^(\\multicolumn%s*(%b{})%s*(%b{}))")
  if head then
    local br = t:find("{", #head + 1, true)
    local close = br and find_close(t, br)
    if close then
      cell.colspan = tonumber(a:sub(2, -2)) or 1
      cell.spec = b:sub(2, -2)
      cell.body = t:sub(br + 1, close - 1)
      t = trim(t:sub(close + 1))
      if t ~= "" then cell.body = cell.body .. " " .. t end
      return cell
    end
  end
  local mhead, mn = t:match("^(\\multirow%s*(%b{}))")
  if mhead then
    -- optional [vpos], then {width}{content}
    local rest = t:sub(#mhead + 1)
    local opt = rest:match("^%b[]")
    if opt then rest = rest:sub(#opt + 1) end
    local w = rest:match("^%s*%b{}")
    if w then
      rest = rest:sub(#rest:match("^%s*") + #w + 1)
      local br = rest:find("{", 1, true)
      local close = br and find_close(rest, br)
      if close then
        cell.rowspan = tonumber(mn:sub(2, -2)) or 1
        cell.body = rest:sub(br + 1, close - 1)
        return cell
      end
    end
  end
  return cell
end

local function parse_body(body)
  local rows, trailing = {}, nil
  local segs = split_top(body, "\\\\")
  for idx, raw in ipairs(segs) do
    local seg, rules = take_rules(raw)
    seg = trim(seg)
    if seg == "" then
      if idx == #segs then
        trailing = rules
      elseif #rows > 0 and has_rule(rules) then
        -- a rule on its own line between rows: attach below previous row
        rows[#rows].below = rules
      end
    else
      local cells = {}
      for _, c in ipairs(split_top(seg, "&")) do
        cells[#cells + 1] = parse_cell(c)
      end
      rows[#rows + 1] = { cells = cells, above = rules }
    end
  end
  return rows, trailing
end

------------------------------------------------------------------ html

local function cell_inlines(txt)
  txt = trim(txt)
  if txt == "" then return pandoc.Inlines({}) end
  local ok, doc = pcall(pandoc.read, txt, "latex")
  if ok and doc and #doc.blocks > 0 then
    return pandoc.utils.blocks_to_inlines(doc.blocks)
  end
  return pandoc.Inlines({ pandoc.Str(txt) })
end

local function rule_class(kind, where)
  if kind == "thick" then return "ltx-" .. where .. "2" end
  if kind == "double" then return "ltx-" .. where .. "d" end
  if kind == "thin" then return "ltx-" .. where .. "1" end
  return nil
end

local ALIGN = { l = "AlignLeft", c = "AlignCenter", r = "AlignRight" }

-- build one pandoc.Row.
-- `pending` maps column -> number of further rows still covered by a
-- \multirow above. LaTeX writes an empty placeholder cell for those, but
-- HTML's rowspan already covers them, so they must be consumed, not emitted.
local function build_row(cols, row, ri, nrows, trailing, pending)
  local cells = pandoc.List({})
  local ci = 1
  local input, k = row.cells, 1
  while k <= #input do
    -- skip columns still covered from above
    while (pending[ci] or 0) > 0 do
      pending[ci] = pending[ci] - 1
      ci = ci + 1
      if trim(input[k].body) == "" and input[k].colspan == 1
          and input[k].rowspan == 1 then
        k = k + 1 -- eat the placeholder
        if k > #input then break end
      end
    end
    if k > #input then break end
    local cell = input[k]
    k = k + 1
    local col = cols[ci] or { align = "l", left = 0, right = 0 }
    local last = cols[math.min(ci + cell.colspan - 1, #cols)] or col
    local align, left, right = col.align, col.left, last.right
    local width = col.width
    if cell.spec ~= nil then
      local sub = parse_colspec(cell.spec)
      if #sub > 0 then
        align = sub[1].align
        left = sub[1].left
        right = sub[#sub].right
      end
    end
    local cls = pandoc.List({ "ltx-" .. align })
    if left > 0 then cls:insert("ltx-bl") end
    if right > 0 then cls:insert("ltx-br") end

    local ra = row.above
    if ra.all then
      cls:insert(rule_class(ra.all, "t"))
    else
      for _, rg in ipairs(ra.ranges) do
        if ci >= rg[1] and ci <= rg[2] then cls:insert("ltx-t1") break end
      end
    end
    local rb = row.below or (ri == nrows and trailing or nil)
    if rb then
      if rb.all then
        cls:insert(rule_class(rb.all, "b"))
      else
        for _, rg in ipairs(rb.ranges) do
          if ci >= rg[1] and ci <= rg[2] then cls:insert("ltx-b1") break end
        end
      end
    end

    local attrs = {}
    if width and width ~= "auto" then attrs.style = "width:" .. width end
    local content = pandoc.Blocks({ pandoc.Plain(cell_inlines(cell.body)) })
    cells:insert(pandoc.Cell(content, ALIGN[align] or "AlignDefault",
      cell.rowspan, cell.colspan, pandoc.Attr("", cls, attrs)))
    if cell.rowspan > 1 then
      for c = ci, ci + cell.colspan - 1 do
        pending[c] = (pending[c] or 0) + cell.rowspan - 1
      end
    end
    ci = ci + cell.colspan
  end
  return pandoc.Row(cells)
end

local function render_table(cols, rows, trailing, header_end, id, caption)
  local colspecs = pandoc.List({})
  for _, c in ipairs(cols) do
    colspecs:insert({ ALIGN[c.align] or "AlignDefault", nil })
  end

  local head_rows, body_rows = pandoc.List({}), pandoc.List({})
  local pending = {}
  for ri, row in ipairs(rows) do
    local r = build_row(cols, row, ri, #rows, trailing, pending)
    if header_end and ri <= header_end then
      head_rows:insert(r)
    else
      body_rows:insert(r)
    end
  end

  local cap = pandoc.Caption({})
  if caption ~= nil then
    cap = pandoc.Caption(pandoc.Blocks({ pandoc.Plain(cell_inlines(caption)) }))
  end

  return pandoc.Table(
    cap,
    colspecs,
    pandoc.TableHead(head_rows),
    { { attr = pandoc.Attr(), body = body_rows,
        head = pandoc.List({}), row_head_columns = 0 } },
    pandoc.TableFoot(),
    pandoc.Attr(id or "", { "ltx-table" })
  )
end

------------------------------------------------------------------ transform

local function find_env(txt, name)
  local st = txt:find("\\begin%s*{" .. name .. "}")
  if st == nil then return nil end
  local _, be = txt:find("\\begin%s*{" .. name .. "}", st)
  local es, ee = txt:find("\\end%s*{" .. name .. "}", be)
  if es == nil then return nil end
  return st, be, es, ee
end

local function tabular_to_table(txt, id, caption)
  local name
  for env in pairs(ENVS) do
    if txt:find("\\begin%s*{" .. env:gsub("%*", "%%*") .. "}") then
      name = env
      break
    end
  end
  if name == nil then return nil end
  local esc_name = name:gsub("%*", "%%*")
  local _, be, es = find_env(txt, esc_name)
  if be == nil then return nil end
  local rest = txt:sub(be + 1)
  local inner_end = es - be - 1
  local inner = rest:sub(1, inner_end)

  -- tabularx/tabular* take a width argument first
  if name == "tabularx" or name == "tabular*" then
    local w = inner:match("^%s*%b{}")
    if w then inner = inner:sub(#inner:match("^%s*") + #w + 1) end
  end
  -- optional [t]/[b] positioning
  local pos = inner:match("^%s*%b[]")
  if pos then inner = inner:sub(#inner:match("^%s*") + #pos + 1) end

  local specm = inner:match("^%s*%b{}")
  if specm == nil then return nil end
  local body = inner:sub(#inner:match("^%s*") + #specm + 1)
  local spec = specm:sub(2, -2)

  local no_header = body:find("no%-header") ~= nil
  body = strip_comments(body)
  local cols = parse_colspec(spec)
  if #cols == 0 then return nil end
  local rows, trailing = parse_body(body)
  if #rows == 0 then return nil end

  local header_end = 0
  if not no_header then
    for i = 2, #rows do
      if has_rule(rows[i].above) or (rows[i - 1].below and has_rule(rows[i - 1].below)) then
        header_end = i - 1
        break
      end
    end
  end
  return render_table(cols, rows, trailing, header_end, id, caption)
end

-- \begin{table} ... \caption{...} \label{...} ... \end{table}
local function float_parts(txt)
  if not txt:find("\\begin%s*{table%*?}") then return nil end
  local cap = nil
  local cs = txt:find("\\caption")
  if cs then
    local br = txt:find("{", cs, true)
    local close = br and find_close(txt, br)
    if close then cap = txt:sub(br + 1, close - 1) end
  end
  local lab = txt:match("\\label%s*{(.-)}")
  return { caption = cap, label = lab }
end

local function RawBlock(el)
  if el.format ~= "tex" and el.format ~= "latex" then return nil end
  local txt = el.text
  local has_tab = false
  for env in pairs(ENVS) do
    if txt:find("\\begin%s*{" .. env:gsub("%*", "%%*") .. "}") then
      has_tab = true
      break
    end
  end
  if not has_tab then return nil end
  if fmt_is("latex") then return nil end -- PDF: it is already LaTeX
  if not fmt_is("html") then return nil end

  local float = float_parts(txt) or {}
  local tbl = tabular_to_table(txt, float.label, float.caption)
  if tbl == nil then
    io.stderr:write("[latex-tables] could not parse a tabular; left as-is\n")
    return nil
  end
  return tbl
end

local function read_meta(meta)
  if quarto ~= nil and quarto.doc ~= nil then
    if fmt_is("html") then
      quarto.doc.add_html_dependency({
        name = "calc-latex-tables",
        version = "1.0.0",
        stylesheets = { "latex-tables.css" },
      })
    elseif fmt_is("latex") then
      quarto.doc.use_latex_package("booktabs")   -- \toprule etc.
      quarto.doc.use_latex_package("multirow")   -- \multirow
      quarto.doc.use_latex_package("array")
    end
  end
  return meta
end

return {
  { Meta = read_meta },
  { RawBlock = RawBlock },
}