-- Lift Quarto's 1x1 docx float wrapper so captions sit outside the
-- body. Tables: caption above, Word style 表题. Figures: caption
-- below, Word style 图题. Both centered.
--
-- Quarto 1.4+ docx FloatRefTarget renderer emits a SimpleTable whose
-- only cell is a Div{caption, body}. Caption Paras start with
-- RawInline openxml that hard-codes ImageCaption and jc=left.
--
-- Docx only. Do not set .attr on Para/Table: Quarto's walker rejects it.
-- Idempotent. Run at post-render after the numbering filter
-- (crossref-tight-num.lua or crossref-book-num.lua).
-- Accepts website 图N- and book 图3-1 / 图3.1 captions.

local TBL = "\232\161\168" -- U+8868 表
local FIG = "\229\155\190" -- U+56FE 图
local STYLE_TBL = TBL .. "\233\162\152" -- 表题
local STYLE_FIG = FIG .. "\233\162\152" -- 图题
local STYLE_ID_TBL = "BiaoTi"
local STYLE_ID_FIG = "TuTi"

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 element_id(el)
  local ok, ident = pcall(function()
    return el.identifier
  end)
  if ok and type(ident) == "string" then
    return ident
  end
  return ""
end

local function first_char(text)
  if not text or text == "" then
    return ""
  end
  return pandoc.text.sub(text, 1, 1)
end

local function is_numbered_float_text(text, prefix)
  if not text or text == "" then
    return false
  end
  local first = first_char(text)
  if first ~= prefix then
    -- literal CJK fallback if file encoding stored the character
    if prefix == TBL and first ~= "表" then
      return false
    end
    if prefix == FIG and first ~= "图" then
      return false
    end
    if prefix ~= TBL and prefix ~= FIG then
      return false
    end
    if first ~= "表" and first ~= "图" then
      return false
    end
  end
  local rest = pandoc.text.sub(text, 2)
  local i = 1
  local n = pandoc.text.len(rest)
  while i <= n do
    local ch = pandoc.text.sub(rest, i, i)
    if ch ~= " " and ch ~= "\u{A0}" and ch ~= "\t" then
      break
    end
    i = i + 1
  end
  if i > 1 then
    rest = pandoc.text.sub(rest, i)
  end
  -- website: 图2- / 图2-13 / 图2- 13 (digits may follow the hyphen after a gap)
  -- book: 图3-1, or 图3.1 / 图3.1a before crossref-book-num.lua rewrites the dot
  if rest:match("^%d+%-") then
    return true
  end
  return rest:match("^%d+%.%d+[A-Za-z]?") ~= nil
end

local function is_tbl_caption_text(text)
  return is_numbered_float_text(text, TBL)
end

local function is_fig_caption_text(text)
  return is_numbered_float_text(text, FIG)
end

local function is_tbl_id(s)
  return type(s) == "string" and s:match("^tbl%-") ~= nil
end

local function is_fig_id(s)
  return type(s) == "string" and s:match("^fig%-") ~= nil
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 wrapper_cell_blocks(tbl)
  if tbl == nil or tbl.t ~= "Table" then
    return nil
  end
  if tbl.colspecs == nil or #tbl.colspecs ~= 1 then
    return nil
  end
  if tbl.bodies == nil or #tbl.bodies ~= 1 then
    return nil
  end
  local rows = tbl.bodies[1].body
  if rows == nil or #rows ~= 1 then
    return nil
  end
  local cells = row_cells(rows[1])
  if cells == nil or #cells ~= 1 then
    return nil
  end
  return cell_blocks(cells[1])
end

local function flatten_divs(blocks, ident)
  ident = ident or ""
  local cur = blocks
  while cur ~= nil and #cur == 1 and cur[1].t == "Div" do
    local did = element_id(cur[1])
    if ident == "" and did ~= "" then
      ident = did
    end
    cur = cur[1].content
  end
  return cur, ident
end

local function has_image(el)
  if el == nil then
    return false
  end
  if el.t == "Image" or el.t == "Figure" then
    return true
  end
  if el.content == nil then
    return false
  end
  local ok, items = pcall(function()
    return el.content
  end)
  if not ok or items == nil then
    return false
  end
  for _, c in ipairs(items) do
    if has_image(c) then
      return true
    end
  end
  return false
end

local function is_body_table(el)
  return el ~= nil and el.t == "Table"
end

local function is_body_figure(el)
  if el == nil then
    return false
  end
  if el.t == "Image" or el.t == "Figure" then
    return true
  end
  if el.t == "Para" or el.t == "Plain" then
    return has_image(el)
  end
  return false
end

local function looks_like_caption_block(el)
  if el == nil then
    return false
  end
  if is_body_table(el) or is_body_figure(el) then
    return false
  end
  local text = pandoc.utils.stringify(el)
  if is_tbl_caption_text(text) or is_fig_caption_text(text) then
    return true
  end
  if el.t == "Para" or el.t == "Plain" then
    for _, inl in ipairs(el.content) do
      if inl.t == "RawInline" and inl.text then
        if inl.text:match("ImageCaption") or inl.text:match("TableCaption")
          or inl.text:match("BiaoTi") or inl.text:match("TuTi") then
          return true
        end
      end
    end
  end
  return false
end

local function classify_blocks(blocks)
  local captions = pandoc.List()
  local bodies = pandoc.List()
  local extras = pandoc.List()
  local kind = nil
  if blocks == nil then
    return nil, captions, bodies, extras
  end
  for _, b in ipairs(blocks) do
    if is_body_table(b) then
      if kind == nil then
        kind = "tbl"
      end
      bodies:insert(b)
    elseif is_body_figure(b) then
      kind = "fig"
      bodies:insert(b)
    elseif looks_like_caption_block(b) then
      captions:insert(b)
    else
      extras:insert(b)
    end
  end
  if kind == nil then
    if #captions > 0 and is_fig_caption_text(pandoc.utils.stringify(captions)) then
      kind = "fig"
    elseif #captions > 0 and is_tbl_caption_text(pandoc.utils.stringify(captions)) then
      kind = "tbl"
    end
  end
  return kind, captions, bodies, extras
end

local function retarget_caption_openxml(inlines, style_id)
  local out = pandoc.List()
  for _, el in ipairs(inlines) do
    if el.t == "RawInline" and el.format and tostring(el.format):match("openxml") then
      local t = el.text or ""
      t = t:gsub("ImageCaption", style_id)
      t = t:gsub("TableCaption", style_id)
      t = t:gsub("BiaoTi", style_id)
      t = t:gsub("TuTi", style_id)
      -- Quarto docxDivCaption hard-codes jc=left; 表题/图题 are center.
      t = t:gsub('w:jc w:val="left"', 'w:jc w:val="center"')
      t = t:gsub("w:jc w:val='left'", "w:jc w:val='center'")
      out:insert(pandoc.RawInline("openxml", t))
    else
      out:insert(el)
    end
  end
  return out
end

local function caption_inlines(block, style_id)
  if block.t == "Para" or block.t == "Plain" then
    return retarget_caption_openxml(block.content, style_id)
  end
  return retarget_caption_openxml(pandoc.utils.blocks_to_inlines({ block }), style_id)
end

local function styled_caption_div(block, ident, style_name, style_id)
  local para = pandoc.Para(caption_inlines(block, style_id))
  local id = ""
  if ident and ident ~= "" then
    id = ident
  end
  return pandoc.Div({ para }, pandoc.Attr(id, {}, { ["custom-style"] = style_name }))
end

local function style_caption_blocks(blocks, ident, style_name, style_id)
  local out = pandoc.List()
  local used_ident = false
  for _, b in ipairs(blocks) do
    local id = nil
    if not used_ident then
      id = ident
      used_ident = true
    end
    out:insert(styled_caption_div(b, id, style_name, style_id))
  end
  return out
end

local function take_ident(wrapper, captions, bodies, div_ident, kind)
  if kind == "fig" then
    if is_fig_id(div_ident) then
      return div_ident
    end
    local wid = element_id(wrapper)
    if is_fig_id(wid) then
      return wid
    end
    for _, b in ipairs(captions) do
      if is_fig_id(element_id(b)) then
        return element_id(b)
      end
    end
    for _, b in ipairs(bodies) do
      if is_fig_id(element_id(b)) then
        return element_id(b)
      end
    end
    if div_ident and div_ident ~= "" then
      return div_ident
    end
    return wid
  end
  if is_tbl_id(div_ident) then
    return div_ident
  end
  local wid = element_id(wrapper)
  if is_tbl_id(wid) then
    return wid
  end
  for _, b in ipairs(captions) do
    if is_tbl_id(element_id(b)) then
      return element_id(b)
    end
  end
  for _, b in ipairs(bodies) do
    if is_tbl_id(element_id(b)) then
      return element_id(b)
    end
  end
  if div_ident and div_ident ~= "" then
    return div_ident
  end
  return wid
end

local function looks_like_kind(captions, ident, kind)
  local text = pandoc.utils.stringify(captions)
  if kind == "fig" then
    if is_tbl_caption_text(text) then
      return false
    end
    return is_fig_caption_text(text) or is_fig_id(ident)
  end
  if is_fig_caption_text(text) then
    return false
  end
  return is_tbl_caption_text(text) or is_tbl_id(ident)
end

local function unwrap_table(tbl)
  local cell = wrapper_cell_blocks(tbl)
  if cell == nil then
    return nil
  end
  local flat, div_ident = flatten_divs(cell, "")
  local kind, captions, bodies, extras = classify_blocks(flat)
  if kind == nil or #captions == 0 or #bodies == 0 then
    return nil
  end
  local ident = take_ident(tbl, captions, bodies, div_ident, kind)
  if not looks_like_kind(captions, ident, kind) then
    return nil
  end
  local style_name = STYLE_TBL
  local style_id = STYLE_ID_TBL
  if kind == "fig" then
    style_name = STYLE_FIG
    style_id = STYLE_ID_FIG
  end
  local styled = style_caption_blocks(captions, ident, style_name, style_id)
  local out = pandoc.List()
  -- tables: caption above; figures: caption below
  if kind == "fig" then
    for _, b in ipairs(bodies) do
      out:insert(b)
    end
    for _, b in ipairs(extras) do
      out:insert(b)
    end
    for _, b in ipairs(styled) do
      out:insert(b)
    end
  else
    for _, b in ipairs(styled) do
      out:insert(b)
    end
    for _, b in ipairs(bodies) do
      out:insert(b)
    end
    for _, b in ipairs(extras) do
      out:insert(b)
    end
  end
  return out
end

local function walk_blocks(blocks)
  local out = pandoc.List()
  for _, b in ipairs(blocks) do
    if b.t == "Table" then
      local repl = unwrap_table(b)
      if repl then
        for _, x in ipairs(repl) do
          out:insert(x)
        end
      else
        out:insert(b)
      end
    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