-- Recompute docx table column widths and bold the header row.
-- Pandoc sizes pipe-table columns by Markdown character count (CJK=1),
-- then writes tblLayout=fixed. Short columns such as 年份/2024 get ~4%
-- width; cell padding eats the rest and digits wrap to several lines.
-- Run at post-render after docx-table-caption.lua. Docx only.
-- Do not set .attr on Table/Cell: Quarto's walker rejects it.

local SHORT_UNITS = 6
local SHORT_MIN_FRAC = 0.11
local PAD_UNITS = 2

local function is_docx()
  if quarto and quarto.doc and quarto.doc.is_format then
    local ok, res = pcall(function()
      return quarto.doc.is_format("docx")
    end)
    if ok and res then
      return true
    end
  end
  return FORMAT ~= nil and tostring(FORMAT):match("docx") ~= nil
end

local function is_cjk(cp)
  return (cp >= 0x2E80 and cp <= 0x9FFF)
    or (cp >= 0xF900 and cp <= 0xFAFF)
    or (cp >= 0xFF00 and cp <= 0xFFEF)
    or (cp >= 0x3000 and cp <= 0x303F)
end

local function display_width(text)
  if not text or text == "" then
    return 0
  end
  local w = 0
  local ok, _ = pcall(function()
    for _, cp in utf8.codes(text) do
      if is_cjk(cp) then
        w = w + 2
      else
        w = w + 1
      end
    end
  end)
  if not ok then
    return pandoc.text.len(text)
  end
  return w
end

local function row_cells(row)
  if row == nil then
    return nil
  end
  if row.cells ~= nil then
    return row.cells
  end
  return row
end

local function cell_blocks(cell)
  if cell == nil then
    return nil
  end
  if cell.contents ~= nil then
    return cell.contents
  end
  return cell
end

local function spec_align(spec)
  if spec == nil then
    return pandoc.AlignCenter
  end
  if spec.alignment ~= nil then
    return spec.alignment
  end
  return spec[1] or pandoc.AlignCenter
end

local function for_each_row(tbl, fn)
  if tbl.head and tbl.head.rows then
    for _, row in ipairs(tbl.head.rows) do
      fn(row)
    end
  end
  if tbl.bodies then
    for _, body in ipairs(tbl.bodies) do
      if body.head then
        for _, row in ipairs(body.head) do
          fn(row)
        end
      end
      if body.body then
        for _, row in ipairs(body.body) do
          fn(row)
        end
      end
    end
  end
  if tbl.foot and tbl.foot.rows then
    for _, row in ipairs(tbl.foot.rows) do
      fn(row)
    end
  end
end

local function already_strong(inlines)
  return inlines ~= nil and #inlines == 1 and inlines[1].t == "Strong"
end

local function strong_blocks(blocks)
  if blocks == nil then
    return blocks
  end
  local out = pandoc.List()
  for _, b in ipairs(blocks) do
    -- Keep Plain: Pandoc's docx writer maps table-cell Plain to Compact (五号).
    -- Para drops Compact and falls back to Normal (小四).
    if (b.t == "Para" or b.t == "Plain") and not already_strong(b.content) then
      out:insert(pandoc.Plain({ pandoc.Strong(b.content) }))
    else
      out:insert(b)
    end
  end
  return out
end

local function bold_header(tbl)
  if tbl.head == nil or tbl.head.rows == nil then
    return
  end
  for _, row in ipairs(tbl.head.rows) do
    local cells = row_cells(row)
    if cells then
      for _, cell in ipairs(cells) do
        local blocks = cell_blocks(cell)
        if blocks ~= nil and cell.contents ~= nil then
          cell.contents = strong_blocks(blocks)
        end
      end
    end
  end
end

local function measure_columns(tbl, ncols)
  local widths = {}
  for i = 1, ncols do
    widths[i] = 0
  end
  for_each_row(tbl, function(row)
    local cells = row_cells(row)
    if cells == nil then
      return
    end
    for i, cell in ipairs(cells) do
      if i <= ncols then
        local w = display_width(pandoc.utils.stringify(cell_blocks(cell) or {}))
        if w > widths[i] then
          widths[i] = w
        end
      end
    end
  end)
  return widths
end

local function fractions_from_widths(raw)
  local n = #raw
  local padded = {}
  for i = 1, n do
    local w = raw[i] + PAD_UNITS
    if w < 4 then
      w = 4
    end
    padded[i] = w
  end
  local sum = 0
  for i = 1, n do
    sum = sum + padded[i]
  end
  if sum <= 0 then
    local even = 1 / n
    local frac = {}
    for i = 1, n do
      frac[i] = even
    end
    return frac
  end
  local frac = {}
  for i = 1, n do
    frac[i] = padded[i] / sum
    if raw[i] <= SHORT_UNITS and frac[i] < SHORT_MIN_FRAC then
      frac[i] = SHORT_MIN_FRAC
    end
  end
  local tot = 0
  for i = 1, n do
    tot = tot + frac[i]
  end
  for i = 1, n do
    frac[i] = frac[i] / tot
  end
  return frac
end

local function apply_colspecs(tbl, frac)
  local specs = pandoc.List()
  for i, spec in ipairs(tbl.colspecs) do
    specs:insert({ spec_align(spec), frac[i] })
  end
  tbl.colspecs = specs
end

local function fix_table(tbl)
  if tbl == nil or tbl.t ~= "Table" or tbl.colspecs == nil then
    return tbl
  end
  local n = #tbl.colspecs
  if n < 2 then
    return tbl
  end
  local raw = measure_columns(tbl, n)
  apply_colspecs(tbl, fractions_from_widths(raw))
  bold_header(tbl)
  return tbl
end

local function walk_blocks(blocks)
  local out = pandoc.List()
  for _, b in ipairs(blocks) do
    if b.t == "Table" then
      out:insert(fix_table(b))
    elseif b.t == "Div" and b.content ~= nil then
      b.content = walk_blocks(b.content)
      out:insert(b)
    else
      out:insert(b)
    end
  end
  return out
end

function Pandoc(doc)
  if not is_docx() then
    return doc
  end
  doc.blocks = walk_blocks(doc.blocks)
  return doc
end