-- calc-notes.lua
-- "calc" / "calc*" derivation environments for Quarto (HTML + PDF).
--
--   \begin{calc}
--   f(x) &= (x+1)^2          \note{expand the square}\\
--        &= x^2 + 2x + 1     \mnote{$(a+b)^2 = a^2+2ab+b^2$}
--   \end{calc}
--
-- \note{...}  -> inline note next to the line   (foldable in HTML, plain colored text in PDF)
-- \mnote{...} -> math note: content is raw math, no $...$ needed; may contain
--                aligned lines like "f &= ma \\\\ &= m\\dot{v}"
-- calc*       -> the last line is boxed to highlight the final result
--
-- Rename the environment here if you want something other than "calc":
local ENV = "calc"
local DEFAULT_COLOR = "#B45309" -- overridden by `notecolor:` in the document metadata

local notecolor = DEFAULT_COLOR

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

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

-- Works both under Quarto and plain pandoc.
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

-- Split `s` on every occurrence of the literal `sep` that sits at brace
-- depth 0 AND outside any nested environment, so neither "\\" inside
-- \text{...} nor the row/column separators of a nested \begin{pmatrix},
-- \begin{cases}, \begin{array}... ever split a calc line.
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 -- skip escaped char / start of a control sequence
        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

-- s:sub(open_idx) == "{"; returns index of the matching "}".
local function find_balanced_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

-- Remove every \cmd{...} from `line`; return (cleaned line, {contents...}).
local function extract_cmd(line, cmd)
  local found = {}
  while true do
    local st, brace = line:find("\\" .. cmd .. "%s*{")
    if not st then break end
    local close = find_balanced_close(line, brace)
    if not close then break end
    found[#found + 1] = line:sub(brace + 1, close - 1)
    line = line:sub(1, st - 1) .. line:sub(close + 1)
  end
  return line, found
end

-- \mnote content is raw math; wrap multi-line/aligned content in aligned.
-- `opt` is the aligned position option ("[t]" for LaTeX, "" for MathJax).
local function wrap_mnote(txt, opt)
  if #split_top(txt, "\\\\") > 1 or #split_top(txt, "&") > 1 then
    return "\\begin{aligned}" .. opt .. " " .. txt .. " \\end{aligned}"
  end
  return txt
end

------------------------------------------------------------------ parse

local function parse_body(body)
  local lines = {}
  for _, raw in ipairs(split_top(body, "\\\\")) do
    local l = trim(raw)
    -- drop an optional spacing argument left over from "\\[1ex]"
    l = trim(l:gsub("^%[[^%]]*%]", ""))
    if l ~= "" then
      local m, notes = extract_cmd(l, "note")
      local m2, mnotes = extract_cmd(m, "mnote")
      lines[#lines + 1] = { math = trim(m2), notes = notes, mnotes = mnotes }
    end
  end
  return lines
end

------------------------------------------------------------------ HTML

-- Note contents are parsed as LaTeX so `$...$` inside a note renders as math.
local function note_inlines(txt)
  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 details_blocks(content, kind)
  local b = pandoc.Blocks({})
  b:insert(pandoc.RawBlock("html",
    '<details class="calc-note calc-note-' .. kind .. '"><summary>note</summary><div class="calc-note-body">'))
  if kind == "math" then
    -- \mnote content is raw math — no $...$ needed
    b:insert(pandoc.Plain {
      pandoc.Math("InlineMath", wrap_mnote(content, ""))
    })
  else
    b:insert(pandoc.Plain(note_inlines(content)))
  end
  b:insert(pandoc.RawBlock("html", "</div></details>"))
  return b
end

local function math_div(cls, tex)
  return pandoc.Div(
    pandoc.Plain { pandoc.Math("InlineMath", "\\displaystyle " .. tex) },
    pandoc.Attr("", { cls }))
end

local function html_env(lines, star)
  local rows = pandoc.Blocks({})
  for idx, L in ipairs(lines) do
    local final = star and idx == #lines
    local cells = pandoc.Blocks({})

    if final then
      -- boxed final result: render the whole line (alignment "&" stripped)
      local flat = trim(table.concat(split_top(L.math, "&"), " "))
      cells:insert(pandoc.Div(
        pandoc.Plain {
          pandoc.RawInline("html", '<span class="calc-box">'),
          pandoc.Math("InlineMath", "\\displaystyle {} " .. flat),
          pandoc.RawInline("html", "</span>"),
        },
        pandoc.Attr("", { "calc-full" })))
    else
      local parts = split_top(L.math, "&")
      if #parts == 1 then
        cells:insert(math_div("calc-full", parts[1]))
      else
        local lhs = trim(parts[1])
        local rhs = trim(table.concat(parts, " ", 2))
        if lhs ~= "" then
          cells:insert(math_div("calc-lhs", lhs))
        end
        -- leading "{}" keeps correct spacing for a line starting with "= ..."
        cells:insert(math_div("calc-rhs", "{} " .. rhs))
      end
    end

    local nb = pandoc.Blocks({})
    for _, t in ipairs(L.notes) do nb:extend(details_blocks(t, "inline")) end
    for _, t in ipairs(L.mnotes) do nb:extend(details_blocks(t, "math")) end
    if #nb > 0 then
      cells:insert(pandoc.Div(nb, pandoc.Attr("", { "calc-notes" })))
    end

    local rowcls = { "calc-row" }
    if final then rowcls[#rowcls + 1] = "calc-final" end
    rows:insert(pandoc.Div(cells, pandoc.Attr("", rowcls)))
  end
  local cls = star and { "calc", "calc-star" } or { "calc" }
  return pandoc.Blocks { pandoc.Div(rows, pandoc.Attr("", cls)) }
end

------------------------------------------------------------------ LaTeX

local function latex_env(lines, star)
  local out = {}
  for idx, L in ipairs(lines) do
    local seg = L.math
    if star and idx == #lines then
      if #split_top(seg, "&") > 1 then
        seg = "\\Aboxed{" .. seg .. "}" -- boxes across the alignment point (mathtools)
      else
        seg = "\\boxed{" .. seg .. "}"
      end
    end
    for _, t in ipairs(L.notes) do
      seg = seg .. " \\quad {\\color{notecolor}\\text{\\small " .. t .. "}}"
    end
    for _, t in ipairs(L.mnotes) do
      -- math note: content is math already; \text{\small$..$} scales it down
      seg = seg .. " \\quad {\\color{notecolor}\\text{\\small$"
        .. wrap_mnote(t, "[t]") .. "$}}"
    end
    out[#out + 1] = "  " .. seg
  end
  return pandoc.Blocks {
    pandoc.RawBlock("latex",
      "\\begin{align*}\n" .. table.concat(out, " \\\\\n") .. "\n\\end{align*}")
  }
end

------------------------------------------------------------------ filter

local function transform(txt)
  local star, body = txt:match(
    "^%s*\\begin{" .. ENV .. "(%*?)}(.-)\\end{" .. ENV .. "%*?}%s*$")
  if body == nil then return nil end
  local lines = parse_body(body)
  if #lines == 0 then return nil end
  local isstar = (star == "*")
  if fmt_is("latex") then
    return latex_env(lines, isstar)
  elseif fmt_is("html") then
    return html_env(lines, isstar)
  end
  return nil
end

function RawBlock(el)
  if el.format ~= "tex" and el.format ~= "latex" then return nil end
  return transform(el.text)
end

-- also catch the env when it's wrapped in $$ ... $$
-- The math may share its paragraph with other text (a line touching the
-- fence, a trailing character, an inline "$$...$$" in a sentence). Split
-- the paragraph: text before/after stays as text, the env becomes blocks.
local function is_ws(it)
  return it.t == "Space" or it.t == "SoftBreak" or it.t == "LineBreak"
      or (it.t == "Str" and it.text:match("^%s*$") ~= nil)
end

local function flush(cur, out, ctor)
  local only_ws = true
  for _, it in ipairs(cur) do
    if not is_ws(it) then only_ws = false; break end
  end
  if not only_ws then out:insert(ctor(cur)) end
end

local function block_with_math(el, ctor)
  local out = pandoc.Blocks({})
  local cur = pandoc.Inlines({})
  local found = false
  for _, it in ipairs(el.content) do
    local res = nil
    if it.t == "Math" and it.mathtype == "DisplayMath" then
      res = transform(it.text)
    end
    if res ~= nil then
      found = true
      flush(cur, out, ctor)
      cur = pandoc.Inlines({})
      out:extend(res)
    else
      cur:insert(it)
    end
  end
  if not found then return nil end
  flush(cur, out, ctor)
  return out
end

function Para(el) return block_with_math(el, pandoc.Para) end
function Plain(el) return block_with_math(el, pandoc.Plain) end

function Meta(m)
  if m.notecolor then
    notecolor = pandoc.utils.stringify(m.notecolor)
  end
  if quarto == nil or quarto.doc == nil then return end -- plain pandoc: skip injection
  if fmt_is("html") then
    quarto.doc.add_html_dependency({
      name = "calc-notes",
      version = "1.0.0",
      stylesheets = { "calc-notes.css" },
      scripts = { "calc-notes.js" },
    })
    quarto.doc.include_text("in-header",
      "<style>:root { --calc-notecolor: " .. notecolor .. "; }</style>")
  elseif fmt_is("latex") then
    quarto.doc.use_latex_package("mathtools")  -- \Aboxed
    quarto.doc.use_latex_package("xcolor")
    local hex = notecolor:gsub("^#", "")
    quarto.doc.include_text("in-header", table.concat({
      "\\definecolor{notecolor}{HTML}{" .. hex .. "}",
    }, "\n"))
  end
end