--[[
count_latex_words.lua

Count the number of "real words" in a LaTeX (.tex) document -- i.e. the
words that would actually appear as visible text if the document were
typeset, excluding preamble, command definitions, commands, math, comments,
labels, citations, includes, and similar non-prose content.

This is a pattern/heuristic-based tool, not a full LaTeX parser, so it will
not be 100% perfect on every document (LaTeX is not a regular language).
It mirrors a Python implementation of the same approach and aims to be a
practical approximation that handles the vast majority of real documents.

USAGE AS A NEOVIM PLUGIN MODULE
--------------------------------
    local wc = require("count_latex_words")

    -- Count the current buffer using the default rules:
    local n = wc.count_buffer(0)

    -- Count the current buffer with a custom config file:
    local config = wc.default_config()
    wc.load_config(config, "/path/to/mycmds.json")
    local n = wc.count_buffer(0, config)

    -- Count an arbitrary file or string:
    local n = wc.count_file("paper.tex", config)
    local n = wc.count_text("\\textbf{hello} world", config)

If loaded inside Neovim, this module also registers a `:LatexWordCount`
user command (optionally takes a path to a JSON config file as its arg).

USAGE AS A STANDALONE SCRIPT
------------------------------
    lua count_latex_words.lua paper.tex
    lua count_latex_words.lua paper.tex mycmds.json

CUSTOM COMMANDS (JSON CONFIG)
------------------------------
Pass a JSON file describing how your own macros should be treated:

    {
      "drop_environments":  ["myfigureenv"],
      "drop_commands":      ["mysecretnote"],
      "keep_arg_commands":  ["myterm"],
      "zero_arg_commands":  ["companyname"],
      "custom_arg_commands": {
        "mytranslation": ["drop", "keep"]
      }
    }

  - drop_environments:   environments whose entire contents are discarded.
  - drop_commands:       commands whose arguments carry no visible words;
                          command + all {..} args are removed.
  - keep_arg_commands:   commands wrapping a single word/phrase to keep;
                          command is stripped, its {..} argument is kept.
  - zero_arg_commands:   commands with no argument that expand to nothing
                          countable; just the command token is removed.
  - custom_arg_commands: commands with MULTIPLE {..} arguments where each
                          argument is individually "keep" or "drop", matched
                          in order, e.g. \mytranslation{fr}{en} ->
                          ["drop", "keep"] keeps only the English text.

Config rules extend the built-in defaults. If a command name is listed
under a new category, it is automatically removed from any other category
to avoid ambiguity (so a config file can also *reclassify* a built-in).
--]]

local M = {}

-- ---------------------------------------------------------------------
-- Default rule sets
-- ---------------------------------------------------------------------

local function to_set(list)
  local s = {}
  for _, v in ipairs(list) do
    s[v] = true
  end
  return s
end

local DEFAULT_DROP_ENVIRONMENTS = {
  "equation", "equation*", "align", "align*", "alignat", "alignat*",
  "gather", "gather*", "multline", "multline*", "eqnarray", "eqnarray*",
  "verbatim", "verbatim*", "Verbatim", "lstlisting", "minted", "algorithm",
  "algorithmic", "tikzpicture", "pgfpicture", "comment",
}

local DEFAULT_DROP_COMMANDS = {
  "label", "ref", "eqref", "pageref", "autoref", "nameref", "vref",
  "cite", "citep", "citet", "citeauthor", "citeyear", "citealp", "citealt",
  "nocite", "bibliography", "bibliographystyle",
  "usepackage", "documentclass", "RequirePackage",
  "input", "include", "includeonly", "includegraphics", "includepdf",
  "newcommand", "renewcommand", "providecommand", "DeclareMathOperator",
  "newenvironment", "renewenvironment", "newcolumntype", "definecolor",
  "colorlet", "setlength", "addtolength", "setcounter", "addtocounter",
  "hspace", "vspace", "linespread", "pagestyle",
  "thispagestyle", "graphicspath", "geometry", "hypersetup",
  "bibinput", "newpage", "clearpage", "cleardoublepage",
  "pagebreak", "linebreak", "nopagebreak", "nolinebreak", "appendix",
}

local DEFAULT_KEEP_ARG_COMMANDS = {
  "textbf", "textit", "emph", "underline", "texttt", "textsc", "textsf",
  "textrm", "textnormal", "text", "mbox", "uline", "uwave", "sout",
  "section", "subsection", "subsubsection", "paragraph", "subparagraph",
  "chapter", "part", "caption", "captionof", "footnote", "footnotetext",
  "title", "author", "date", "thanks",
}

local DEFAULT_ZERO_ARG_COMMANDS = {
  "maketitle", "tableofcontents", "listoffigures", "listoftables",
  "noindent", "indent", "par", "newline", "smallskip", "medskip",
  "bigskip", "today", "LaTeX", "LaTeXe", "TeX", "item",
}

-- Commands with MULTIPLE {..} arguments, each individually "keep"/"drop".
local DEFAULT_CUSTOM_ARG_COMMANDS = {
  textcolor = { "drop", "keep" },
  colorbox = { "drop", "keep" },
  href = { "drop", "keep" },
}

function M.default_config()
  local custom = {}
  for k, v in pairs(DEFAULT_CUSTOM_ARG_COMMANDS) do
    custom[k] = { v[1], v[2] }
  end
  return {
    drop_environments = to_set(DEFAULT_DROP_ENVIRONMENTS),
    drop_commands = to_set(DEFAULT_DROP_COMMANDS),
    keep_arg_commands = to_set(DEFAULT_KEEP_ARG_COMMANDS),
    zero_arg_commands = to_set(DEFAULT_ZERO_ARG_COMMANDS),
    custom_arg_commands = custom,
  }
end

function M.build_example_config()
  return {
    drop_environments = { "myfigureenv" },
    drop_commands = { "mysecretnote", "internalcomment" },
    keep_arg_commands = { "myterm", "highlightphrase" },
    zero_arg_commands = { "companyname" },
    custom_arg_commands = {
      mytooltip = { "keep", "drop" },
      mytranslation = { "drop", "keep" },
    },
  }
end

-- ---------------------------------------------------------------------
-- Minimal JSON decoder (fallback for use outside Neovim; Neovim's
-- built-in vim.json.decode / vim.fn.json_decode is preferred when present)
-- ---------------------------------------------------------------------

local json_parse_value -- forward declaration

local function json_skip_ws(s, i)
  while i <= #s and s:sub(i, i):match("%s") do
    i = i + 1
  end
  return i
end

local function json_parse_string(s, i)
  local j = i + 1
  local buf = {}
  local escmap = { ['"'] = '"', ['\\'] = '\\', ['/'] = '/', b = '\b', f = '\f', n = '\n', r = '\r', t = '\t' }
  while j <= #s do
    local c = s:sub(j, j)
    if c == '"' then
      return table.concat(buf), j + 1
    elseif c == "\\" then
      local nc = s:sub(j + 1, j + 1)
      if escmap[nc] then
        table.insert(buf, escmap[nc])
        j = j + 2
      elseif nc == "u" then
        local hex = s:sub(j + 2, j + 5)
        local code = tonumber(hex, 16) or 63
        if code < 0x80 then
          table.insert(buf, string.char(code))
        elseif utf8 and utf8.char then
          table.insert(buf, utf8.char(code))
        else
          table.insert(buf, "?")
        end
        j = j + 6
      else
        table.insert(buf, nc)
        j = j + 2
      end
    else
      table.insert(buf, c)
      j = j + 1
    end
  end
  error("count_latex_words: unterminated string in JSON config")
end

local function json_parse_number(s, i)
  local j = i
  while j <= #s and s:sub(j, j):match("[%d%+%-%.eE]") do
    j = j + 1
  end
  return tonumber(s:sub(i, j - 1)), j
end

local function json_parse_array(s, i)
  local arr = {}
  i = json_skip_ws(s, i + 1)
  if s:sub(i, i) == "]" then
    return arr, i + 1
  end
  while true do
    local val
    val, i = json_parse_value(s, i)
    table.insert(arr, val)
    i = json_skip_ws(s, i)
    local c = s:sub(i, i)
    if c == "," then
      i = json_skip_ws(s, i + 1)
    elseif c == "]" then
      return arr, i + 1
    else
      error("count_latex_words: malformed JSON array in config")
    end
  end
end

local function json_parse_object(s, i)
  local obj = {}
  i = json_skip_ws(s, i + 1)
  if s:sub(i, i) == "}" then
    return obj, i + 1
  end
  while true do
    i = json_skip_ws(s, i)
    if s:sub(i, i) ~= '"' then
      error("count_latex_words: expected string key in JSON config")
    end
    local key
    key, i = json_parse_string(s, i)
    i = json_skip_ws(s, i)
    if s:sub(i, i) ~= ":" then
      error("count_latex_words: expected ':' in JSON config")
    end
    i = json_skip_ws(s, i + 1)
    local val
    val, i = json_parse_value(s, i)
    obj[key] = val
    i = json_skip_ws(s, i)
    local c = s:sub(i, i)
    if c == "," then
      i = i + 1
    elseif c == "}" then
      return obj, i + 1
    else
      error("count_latex_words: malformed JSON object in config")
    end
  end
end

json_parse_value = function(s, i)
  i = json_skip_ws(s, i)
  local c = s:sub(i, i)
  if c == '"' then
    return json_parse_string(s, i)
  elseif c == "{" then
    return json_parse_object(s, i)
  elseif c == "[" then
    return json_parse_array(s, i)
  elseif s:sub(i, i + 3) == "true" then
    return true, i + 4
  elseif s:sub(i, i + 4) == "false" then
    return false, i + 5
  elseif s:sub(i, i + 3) == "null" then
    return nil, i + 4
  else
    return json_parse_number(s, i)
  end
end

function M.json_decode(str)
  if vim ~= nil then
    if vim.json and vim.json.decode then
      local ok, result = pcall(vim.json.decode, str)
      if ok then
        return result
      end
    end
    if vim.fn and vim.fn.json_decode then
      local ok, result = pcall(vim.fn.json_decode, str)
      if ok then
        return result
      end
    end
  end
  local val = json_parse_value(str, 1)
  return val
end

-- ---------------------------------------------------------------------
-- Config loading
-- ---------------------------------------------------------------------

--- Merge rules from a JSON config file into an existing config table.
-- @param config table returned by M.default_config()
-- @param path string path to a JSON config file
function M.load_config(config, path)
  local f, err = io.open(path, "r")
  if not f then
    error("count_latex_words: could not open config file: " .. tostring(err))
  end
  local content = f:read("*a")
  f:close()

  local data = M.json_decode(content)
  if type(data) ~= "table" then
    return config
  end

  local function merge_list(key, target)
    if data[key] then
      for _, name in ipairs(data[key]) do
        target[name] = true
      end
    end
  end

  merge_list("drop_environments", config.drop_environments)
  merge_list("drop_commands", config.drop_commands)
  merge_list("keep_arg_commands", config.keep_arg_commands)
  merge_list("zero_arg_commands", config.zero_arg_commands)

  if data.custom_arg_commands then
    for name, pattern in pairs(data.custom_arg_commands) do
      config.custom_arg_commands[name] = pattern
    end
  end

  -- Reclassify: a command named under a new category is removed from any
  -- other category so config-file rules take priority over defaults.
  local plain_keys = { "drop_commands", "keep_arg_commands", "zero_arg_commands" }
  for _, key in ipairs(plain_keys) do
    if data[key] then
      for _, name in ipairs(data[key]) do
        for _, other in ipairs(plain_keys) do
          if other ~= key then
            config[other][name] = nil
          end
        end
        config.custom_arg_commands[name] = nil
      end
    end
  end
  if data.custom_arg_commands then
    for name, _ in pairs(data.custom_arg_commands) do
      config.drop_commands[name] = nil
      config.keep_arg_commands[name] = nil
      config.zero_arg_commands[name] = nil
    end
  end

  return config
end

-- ---------------------------------------------------------------------
-- Cleaning pipeline
-- ---------------------------------------------------------------------

--- Remove LaTeX comments (unescaped % to end of line).
function M.strip_comments(text)
  local out_lines = {}
  for line in (text .. "\n"):gmatch("(.-)\n") do
    local buf = {}
    local i, n = 1, #line
    while i <= n do
      local c = line:sub(i, i)
      if c == "\\" then
        table.insert(buf, c)
        if i < n then
          table.insert(buf, line:sub(i + 1, i + 1))
          i = i + 2
        else
          i = i + 1
        end
      elseif c == "%" then
        break
      else
        table.insert(buf, c)
        i = i + 1
      end
    end
    table.insert(out_lines, table.concat(buf))
  end
  -- (.."\n"):gmatch("(.-)\n") yields one trailing empty line; drop it if
  -- the original text didn't end with a newline itself is not needed here
  -- since callers only care about the joined result below.
  return table.concat(out_lines, "\n")
end

--- Keep only the content between \begin{document} and \end{document}.
function M.extract_body(text)
  local _, _, body = text:find("\\begin{document}(.-)\\end{document}")
  return body or text
end

--- Escape Lua pattern magic characters in a literal string.
local function escape_pattern(s)
  return (s:gsub("[%(%)%.%%%+%-%*%?%[%]%^%$]", "%%%1"))
end
M.escape_pattern = escape_pattern

--- Remove the entire contents of math/code/graphics-language environments.
function M.drop_environments(text, config)
  for env, _ in pairs(config.drop_environments) do
    local pat = "\\begin{" .. escape_pattern(env) .. "}.-\\end{" .. escape_pattern(env) .. "}"
    text = text:gsub(pat, " ")
  end
  return text
end

-- Lua patterns have no lookbehind, so a plain gsub can't tell an escaped
-- "\$" apart from a real math delimiter. Scan manually instead, treating
-- "$$" as display math and a lone "$" as inline math, skipping either kind
-- when immediately preceded by a backslash.
local function strip_dollar_math(text)
  local out = {}
  local i, n = 1, #text
  while i <= n do
    local c = text:sub(i, i)
    if c == "$" and text:sub(i - 1, i - 1) ~= "\\" then
      local double = text:sub(i + 1, i + 1) == "$"
      local delim = double and "$$" or "$"
      local dlen = double and 2 or 1
      local j = i + dlen
      local found = nil
      while j <= n do
        if text:sub(j, j + dlen - 1) == delim and text:sub(j - 1, j - 1) ~= "\\" then
          found = j
          break
        end
        j = j + 1
      end
      if found then
        table.insert(out, " ")
        i = found + dlen
      else
        -- No closing delimiter found; treat as literal and move on.
        table.insert(out, c)
        i = i + 1
      end
    else
      table.insert(out, c)
      i = i + 1
    end
  end
  return table.concat(out)
end

--- Remove inline/display math: $...$, $$...$$, \(...\), \[...\].
function M.drop_math(text)
  text = strip_dollar_math(text)
  text = text:gsub("\\%(.-\\%)", " ")
  text = text:gsub("\\%[.-\\%]", " ")
  return text
end

local ESCAPED_CHARS = {
  { "\\%%", "%%" },
  { "\\&", "&" },
  { "\\%$", "$" },
  { "\\#", "#" },
  { "\\_", "_" },
  { "\\{", "{" },
  { "\\}", "}" },
}

function M.replace_escaped_chars(text)
  for _, pair in ipairs(ESCAPED_CHARS) do
    text = text:gsub(pair[1], pair[2])
  end
  return text
end

--- If text[pos] starts an optional [..] argument, return the position just
-- past its matching closing bracket; otherwise return pos unchanged.
function M.strip_optional_arg(text, pos)
  if text:sub(pos, pos) == "[" then
    local s, e = text:find("^%b[]", pos)
    if s then
      return e + 1
    end
  end
  return pos
end

--- Given text:sub(pos,pos) == "{", return (content, end_index) where
-- end_index is the index just past the matching closing brace. Handles
-- nesting via Lua's %b{} balanced-match pattern item.
function M.read_braced_group(text, pos)
  local s, e = text:find("^%b{}", pos)
  if s then
    return text:sub(s + 1, e - 1), e + 1
  end
  -- Unbalanced braces -- bail out, consume to end of string.
  return text:sub(pos + 1), #text + 1
end

--- Walk through text once, handling \commandname[opt]{arg}{arg2}... per the
-- config's rule sets, with a fallback for unknown commands that keeps the
-- first {..} argument's text (most custom macros just wrap visible text).
function M.process_commands(text, config)
  local out = {}
  local i, n = 1, #text

  while i <= n do
    local next_bs = text:find("\\", i, true)
    if not next_bs then
      table.insert(out, text:sub(i))
      break
    end
    if next_bs > i then
      table.insert(out, text:sub(i, next_bs - 1))
    end
    i = next_bs

    local s, e, name = text:find("^\\(%a+)%*?", i)
    if s then
      local j = e + 1

      if config.custom_arg_commands[name] then
        j = M.strip_optional_arg(text, j)
        for _, role in ipairs(config.custom_arg_commands[name]) do
          if text:sub(j, j) == "{" then
            local content
            content, j = M.read_braced_group(text, j)
            if role == "keep" then
              table.insert(out, M.process_commands(content, config))
            end
          else
            break
          end
        end
        i = j
      elseif config.drop_commands[name] then
        j = M.strip_optional_arg(text, j)
        while text:sub(j, j) == "{" do
          local _, jn = M.read_braced_group(text, j)
          j = jn
        end
        i = j
      elseif config.zero_arg_commands[name] then
        j = M.strip_optional_arg(text, j)
        i = j
      elseif config.keep_arg_commands[name] then
        j = M.strip_optional_arg(text, j)
        if text:sub(j, j) == "{" then
          local content
          content, j = M.read_braced_group(text, j)
          table.insert(out, M.process_commands(content, config))
        end
        i = j
      elseif name == "begin" or name == "end" then
        j = M.strip_optional_arg(text, j)
        if text:sub(j, j) == "{" then
          local _, jn = M.read_braced_group(text, j)
          j = jn
        end
        i = j
      else
        -- Unknown command: drop the name, skip an optional [..], keep the
        -- first {..} argument's text.
        j = M.strip_optional_arg(text, j)
        if text:sub(j, j) == "{" then
          local content
          content, j = M.read_braced_group(text, j)
          table.insert(out, M.process_commands(content, config))
        end
        i = j
      end
    else
      -- Backslash not followed by a letter: "\\" line break, or stray
      -- backslash (already-escaped chars like \% were handled earlier).
      if text:sub(i, i + 1) == "\\\\" then
        table.insert(out, " ")
        i = i + 2
      else
        i = i + 1
      end
    end
  end

  return table.concat(out)
end

--- Remove stray braces, alignment/table markup, and collapse whitespace.
function M.clean_leftovers(text)
  text = text:gsub("{", " ")
  text = text:gsub("}", " ")
  text = text:gsub("&", " ")
  text = text:gsub("~", " ")
  text = text:gsub("%-%-%-?", " ") -- en/em dashes (-- or ---) as separators
  text = text:gsub("%s+", " ")
  text = text:match("^%s*(.-)%s*$")
  return text
end

--- Tokenize cleaned text into words. A "word" must contain at least one
-- letter or digit (filters out stray punctuation / leftover markup).
function M.count_words(text)
  local words = {}
  for tok in text:gmatch("%S+") do
    local trimmed = tok:gsub("^[%.,;:!%?%(%)%[%]\"'`]+", "")
    trimmed = trimmed:gsub("[%.,;:!%?%(%)%[%]\"'`]+$", "")
    if trimmed:match("%w") then
      table.insert(words, trimmed)
    end
  end
  return words
end

-- ---------------------------------------------------------------------
-- Top-level entry points
-- ---------------------------------------------------------------------

--- Run the full pipeline over a raw LaTeX string.
-- @return count, words (table), cleaned_text
function M.count_text(text, config)
  config = config or M.default_config()
  local t = text
  t = M.strip_comments(t)
  t = M.extract_body(t)
  t = M.drop_environments(t, config)
  t = M.drop_math(t)
  t = M.replace_escaped_chars(t)
  t = M.process_commands(t, config)
  t = M.clean_leftovers(t)
  local words = M.count_words(t)
  return #words, words, t
end

--- Run the full pipeline over a file on disk.
function M.count_file(path, config)
  local f, err = io.open(path, "r")
  if not f then
    error("count_latex_words: could not open file: " .. tostring(err))
  end
  local content = f:read("*a")
  f:close()
  return M.count_text(content, config)
end

--- Run the full pipeline over a Neovim buffer's contents.
-- @param bufnr number|nil buffer handle (0 or nil = current buffer)
function M.count_buffer(bufnr, config)
  if vim == nil then
    error("count_latex_words: count_buffer() requires Neovim's 'vim' API")
  end
  bufnr = bufnr or 0
  local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
  local text = table.concat(lines, "\n")
  return M.count_text(text, config)
end

-- ---------------------------------------------------------------------
-- Neovim plugin wiring (only active when loaded inside Neovim)
-- ---------------------------------------------------------------------

if vim ~= nil and vim.api ~= nil and vim.api.nvim_create_user_command ~= nil then
  vim.api.nvim_create_user_command("LatexWordCount", function(opts)
    local config = M.default_config()
    if opts.args and opts.args ~= "" then
      local ok, err = pcall(M.load_config, config, opts.args)
      if not ok then
        vim.notify(tostring(err), vim.log.levels.ERROR)
        return
      end
    end
    local count = M.count_buffer(0, config)
    vim.notify(string.format("Real words: %d", count))
  end, { nargs = "?", complete = "file", desc = "Count real (visible) words in the current LaTeX buffer" })
end

-- ---------------------------------------------------------------------
-- Standalone CLI usage: `lua count_latex_words.lua paper.tex [config.json]`
-- ---------------------------------------------------------------------

if vim == nil and arg ~= nil and arg[0] and arg[1] then
  local config = M.default_config()
  if arg[2] then
    M.load_config(config, arg[2])
  end
  local count, _, cleaned = M.count_file(arg[1], config)
  if os.getenv("LATEX_WORDS_VERBOSE") then
    io.stderr:write("--- Cleaned text ---\n" .. cleaned .. "\n--- End cleaned text ---\n\n")
  end
  print(string.format("%s: %d words", arg[1], count))
end

return M