--[[
  Pandoc Lua filter: filter_docast.lua
  Converts LaTeX processed by Pandoc into DocAST JSON.

  Usage:
    pandoc --from=latex+raw_tex+latex_macros --to=plain \
           --lua-filter=filter_docast.lua main.tex > build/docast.json

  Notes:
    - Emits a single DocAST JSON object to stdout and exits.
    - Math blocks get stable refs "m/<n>" and are indexed in meta.math_index.
    - Section hierarchy is reconstructed from Header levels.
    - Common theorem-like environments mapped from Div classes: theorem, lemma, corollary, remark, proof.
    - Labels are taken from identifiers (e.g., \label -> {#id} via Pandoc parsing when available).
    - This filter does not expand macros; it only preserves raw TeX for math payloads.
]]

local json = require("pandoc.json")

-- -------- helpers --------

local function now_iso8601()
  -- best-effort local timestamp
  return os.date("!%Y-%m-%dT%H:%M:%SZ")
end

local function deepcopy(t)
  if type(t) ~= "table" then return t end
  local r = {}
  for k,v in pairs(t) do r[k] = deepcopy(v) end
  return r
end

-- stringify (lossless-ish for plain text; math stays separate)
local function stringify(inlines)
  return pandoc.utils.stringify(inlines or {})
end

-- detect a Paragraph that is exactly one DisplayMath inline
local function para_is_single_display_math(para)
  if para.t ~= "Para" then return false end
  local inlines = para.content
  if #inlines ~= 1 then return false end
  local x = inlines[1]
  return (x.t == "Math" and x.mathtype == "DisplayMath")
end

-- extract all inline math from an Inline list; returns:
--   text_inlines_without_math, math_refs[{kind="MathInline", math_ref=...}], created_math_nodes[]
local function extract_inline_math(state, inlines)
  local kept = {}
  local math_child_refs = {}

  for _,el in ipairs(inlines) do
    if el.t == "Math" and el.mathtype == "InlineMath" then
      local ref = state:add_math(el.text, false)
      table.insert(math_child_refs, { kind = "MathInline", math_ref = ref })
      -- replace inline with a lightweight marker in kept text, if desired; omit to keep text clean
      table.insert(kept, pandoc.Str(""))
    else
      table.insert(kept, deepcopy(el))
    end
  end

  return kept, math_child_refs
end

-- map Pandoc Div classes to DocAST "env" kinds
local ENV_MAP = {
  theorem   = "Theorem",
  lemma     = "Lemma",
  corollary = "Corollary",
  remark    = "Remark",
  proof     = "Proof",
  definition= "Definition",
  example   = "Example",
}

-- -------- DocAST builder --------

local Builder = {}
Builder.__index = Builder

function Builder:new(meta)
  local o = {
    root = { uid = "doc/0", kind = "Doc", children = {}, meta = {
      producer = "astx 0.1.0",
      specs = { docast = "1.2.0" },
      generated_utc = now_iso8601(),
      source = { title = pandoc.utils.stringify(meta.title or ""), author = pandoc.utils.stringify(meta.author or "") },
      math_index = {},   -- array of {ref="m/1", display=true|false, tex="..."}
    }},
    stack = {},                 -- section stack
    next_uid = 1,
    next_math = 1,
  }
  setmetatable(o, self)
  return o
end

function Builder:gen_uid(prefix)
  local id = prefix .. "/" .. self.next_uid
  self.next_uid = self.next_uid + 1
  return id
end

function Builder:add_child(parent, node)
  parent.children = parent.children or {}
  table.insert(parent.children, node)
end

function Builder:current_container()
  if #self.stack == 0 then return self.root end
  return self.stack[#self.stack].node
end

function Builder:push_section(level, title, id)
  local node = {
    uid = id ~= "" and id or self:gen_uid("sec"),
    kind = "Section",
    level = level,
    title = title,
    children = {},
  }

  -- find where to attach based on level
  while #self.stack > 0 and self.stack[#self.stack].level >= level do
    table.remove(self.stack)
  end
  local parent = self:current_container()
  self:add_child(parent, node)
  table.insert(self.stack, { level = level, node = node })
end

function Builder:add_para(block)
  local text_inlines, math_refs = extract_inline_math(self, block.content)
  local node = {
    uid = self:gen_uid("p"),
    kind = "Paragraph",
    text = stringify(text_inlines),
    children = math_refs,
  }
  self:add_child(self:current_container(), node)
end

function Builder:add_display_math(tex)
  local ref = self:add_math(tex, true)
  local node = {
    uid = self:gen_uid("eq"),
    kind = "MathBlock",
    math_ref = ref,
  }
  self:add_child(self:current_container(), node)
end

function Builder:add_div_as_env(div)
  -- Determine environment kind from classes
  local kind = nil
  for _,cls in ipairs(div.classes or {}) do
    local k = ENV_MAP[cls]
    if k then kind = k; break end
  end
  if not kind then
    -- pass-through container Div -> DocAST Figure/Paragraph container
    self:add_generic_div(div)
    return
  end

  local node = {
    uid = div.identifier ~= "" and div.identifier or self:gen_uid(string.lower(kind)),
    kind = kind,
    env = string.lower(kind),
    label = div.identifier ~= "" and div.identifier or nil,
    children = {},
  }

  -- Attach to current container
  self:add_child(self:current_container(), node)

  -- Temporarily make it current to collect content
  table.insert(self.stack, { level = 99, node = node }) -- sentinel high level for nesting
  for _,b in ipairs(div.content or {}) do
    self:dispatch_block(b)
  end
  table.remove(self.stack)
end

function Builder:add_generic_div(div)
  local node = {
    uid = div.identifier ~= "" and div.identifier or self:gen_uid("div"),
    kind = "Figure",  -- neutral container; could refine by attributes later
    attrs = { classes = div.classes, attributes = div.attributes },
    children = {},
  }
  self:add_child(self:current_container(), node)
  table.insert(self.stack, { level = 98, node = node })
  for _,b in ipairs(div.content or {}) do
    self:dispatch_block(b)
  end
  table.remove(self.stack)
end

function Builder:add_plain(block)
  local text_inlines, math_refs = extract_inline_math(self, block.content)
  local node = {
    uid = self:gen_uid("plain"),
    kind = "Paragraph",
    text = stringify(text_inlines),
    children = math_refs,
  }
  self:add_child(self:current_container(), node)
end

function Builder:add_blockquote(bq)
  local node = {
    uid = self:gen_uid("quote"),
    kind = "Remark",
    children = {},
  }
  self:add_child(self:current_container(), node)
  table.insert(self.stack, { level = 97, node = node })
  for _,b in ipairs(bq.content or {}) do
    self:dispatch_block(b)
  end
  table.remove(self.stack)
end

function Builder:add_list(list, tight, ordered)
  local node = {
    uid = self:gen_uid(ordered and "ol" or "ul"),
    kind = ordered and "OrderedList" or "BulletList",
    tight = tight or false,
    items = {},
  }
  for _,item in ipairs(list) do
    local itemNode = { uid = self:gen_uid("li"), kind = "ListItem", children = {} }
    table.insert(node.items, itemNode)
    table.insert(self.stack, { level = 96, node = itemNode })
    for _,b in ipairs(item) do self:dispatch_block(b) end
    table.remove(self.stack)
  end
  self:add_child(self:current_container(), node)
end

function Builder:add_table(tbl)
  local node = {
    uid = self:gen_uid("tbl"),
    kind = "Table",
    caption = stringify(tbl.caption and tbl.caption.long or {}),
    children = {}, -- omit cell-level detail for now
  }
  self:add_child(self:current_container(), node)
end

function Builder:add_rule()
  local node = { uid = self:gen_uid("hr"), kind = "ThematicBreak" }
  self:add_child(self:current_container(), node)
end

function Builder:add_codeblock(cb)
  local node = {
    uid = self:gen_uid("code"),
    kind = "CodeBlock",
    language = cb.classes and cb.classes[1] or "",
    attrs = cb.attributes or {},
    text = cb.text or "",
  }
  self:add_child(self:current_container(), node)
end

function Builder:add_rawblock(rb)
  local node = {
    uid = self:gen_uid("raw"),
    kind = "RawBlock",
    format = rb.format,
    text = rb.text,
  }
  self:add_child(self:current_container(), node)
end

function Builder:add_math(self_tex, is_display)
  local ref = "m/" .. tostring(self.next_math)
  self.next_math = self.next_math + 1
  table.insert(self.root.meta.math_index, { ref = ref, display = is_display and true or false, tex = self_tex })
  return ref
end

function Builder:dispatch_block(b)
  local t = b.t
  if t == "Header" then
    local level = b.level
    local title = stringify(b.content)
    local id = b.identifier or ""
    self:push_section(level, title, id)
  elseif t == "Para" then
    if para_is_single_display_math(b) then
      local m = b.content[1]
      self:add_display_math(m.text)
    else
      self:add_para(b)
    end
  elseif t == "Plain" then
    self:add_plain(b)
  elseif t == "Div" then
    self:add_div_as_env(b)
  elseif t == "BlockQuote" then
    self:add_blockquote(b)
  elseif t == "BulletList" then
    self:add_list(b.content or {}, b.tight, false)
  elseif t == "OrderedList" then
    local items = b.content or {}
    self:add_list(items, b.tight, true)
  elseif t == "Table" then
    self:add_table(b)
  elseif t == "HorizontalRule" then
    self:add_rule()
  elseif t == "CodeBlock" then
    self:add_codeblock(b)
  elseif t == "RawBlock" then
    self:add_rawblock(b)
  else
    -- fallback: serialize as generic block container
    local node = {
      uid = self:gen_uid("blk"),
      kind = "Block",
      pandoc_tag = t,
    }
    self:add_child(self:current_container(), node)
  end
end

-- -------- Pandoc entrypoint --------

function Doc(body, meta)
  local B = Builder:new(meta or {})

  -- top-level walk
  for _,b in ipairs(body.blocks or {}) do
    B:dispatch_block(b)
  end

  -- finalize meta: include dialect list placeholder and content hash placeholders
  B.root.meta.dialect = { "nt", "op" }
  B.root.meta.hash = { tex = "", ast = "" }

  -- emit DocAST JSON and exit
  io.write(json.encode(B.root))
  os.exit(0)
end
