path: root/lua/latex-wordcount.lua

--[[
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).

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).

RECURSIVE \input AND \include
------------------------------
count_file() and count_buffer() follow \input{file} and \include{file}
(including the old brace-less `\input file` form) by default, reading
each referenced file, recursing into it (so files can nest arbitrarily
deep), and splicing its content in before counting -- so a main.tex that
pulls in chapter1.tex, chapter2.tex, etc. gets one combined word count.

  - Paths are resolved relative to the ROOT document's directory (exactly
    how real LaTeX/pdflatex resolves \input and \include -- always
    relative to wherever the compiler is invoked, not relative to
    whichever file happens to contain the \input), and ".tex" is appended
    automatically if the bare filename doesn't exist.
  - Already-visited files are skipped (cycle guard), and recursion stops
    after 20 levels deep by default.
  - Missing/unreadable files are left alone rather than aborting the
    whole count; the count_file/count_buffer functions return a 4th
    value, a list of warning strings, so callers can surface problems.
  - Pass { follow_includes = false } as the third argument to disable
    this and count only the file/buffer's own literal content:
        local n = wc.count_file("paper.tex", config, { follow_includes = false })
  - You can also call wc.expand_includes(text, opts) directly if you
    need to splice includes into an arbitrary string yourself.

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,
    do_not_expand = false,
  }
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

-- ---------------------------------------------------------------------
-- Recursive \input / \include expansion
-- ---------------------------------------------------------------------

--- Try to locate a file referenced by \input{name} or \include{name},
-- resolved relative to base_dir, trying the name as given and then with
-- a ".tex" extension appended.
-- @return resolved path string, or nil if no candidate could be opened
function M.resolve_tex_path(filename, base_dir)
  filename = filename:match("^%s*(.-)%s*$") -- trim whitespace
  if filename == "" then
    return nil
  end
  local is_absolute = filename:sub(1, 1) == "/" or filename:match("^%a:[/\\]")
  local candidates
  if is_absolute then
    candidates = { filename, filename .. ".tex" }
  else
    base_dir = base_dir or "."
    candidates = { base_dir .. "/" .. filename, base_dir .. "/" .. filename .. ".tex" }
  end
  for _, c in ipairs(candidates) do
    local f = io.open(c, "r")
    if f then
      f:close()
      return c
    end
  end
  return nil
end

--- Recursively expand \input{file}, \input file (brace-less form), and
-- \include{file} directives by splicing in the referenced file's content.
-- Comments are stripped from each chunk as it's read, so a commented-out
-- \input is correctly left unexpanded.
--
-- @param text string LaTeX source to expand
-- @param opts table|nil:
--   base_dir       directory \input/\include paths are resolved against
--                  (default ".")
--   seen           table used as a set of already-visited resolved paths,
--                  to guard against include cycles (default {})
--   depth          current recursion depth, internal use (default 0)
--   max_depth      recursion limit (default 20)
--   warnings       table that unresolved/cyclic includes get appended to
--                  as human-readable strings (default {})
-- @return expanded text
function M.expand_includes(text, opts)
  opts = opts or {}
  local base_dir = opts.base_dir or "."
  local seen = opts.seen or {}
  local depth = opts.depth or 0
  local max_depth = opts.max_depth or 20
  local warnings = opts.warnings or {}

  text = M.strip_comments(text)

  if depth > max_depth then
    table.insert(warnings, "max include depth (" .. max_depth .. ") exceeded near " .. base_dir)
    return text
  end

  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 matched_name = nil
    for _, name in ipairs({ "input", "include" }) do
      if text:sub(i + 1, i + #name) == name and not text:sub(i + 1 + #name, i + 1 + #name):match("%a") then
        matched_name = name
        break
      end
    end

    if matched_name then
      local j = i + 1 + #matched_name
      j = M.strip_optional_arg(text, j)
      local filename

      if text:sub(j, j) == "{" then
        filename, j = M.read_braced_group(text, j)
      else
        -- Old-style brace-less form: \input filename, terminated by
        -- whitespace, a backslash, or a brace.
        while text:sub(j, j):match("%s") do
          j = j + 1
        end
        local start = j
        while j <= n and not text:sub(j, j):match("[%s\\{}]") do
          j = j + 1
        end
        filename = text:sub(start, j - 1)
      end

      local resolved = M.resolve_tex_path(filename, base_dir)
      if not resolved then
        table.insert(warnings, "could not find included file '" .. filename .. "' (referenced from " .. base_dir .. ")")
      elseif seen[resolved] then
        table.insert(warnings, "skipped already-included file (cycle guard): " .. resolved)
      else
        local f = io.open(resolved, "r")
        if f then
          local included_raw = f:read("*a")
          f:close()
          seen[resolved] = true
          -- Note: real LaTeX resolves \input/\include paths relative to
          -- the ROOT document's directory (i.e. wherever latex/pdflatex
          -- is invoked), not relative to the file doing the including.
          -- So base_dir is deliberately kept the same at every recursion
          -- level, rather than switched to the included file's own dir.
          local expanded = M.expand_includes(included_raw, {
            base_dir = base_dir,
            seen = seen,
            depth = depth + 1,
            max_depth = max_depth,
            warnings = warnings,
          })
          table.insert(out, "\n" .. expanded .. "\n")
        else
          table.insert(warnings, "could not open included file: " .. resolved)
        end
      end

      i = j
    else
      table.insert(out, "\\")
      i = i + 1
    end
  end

  return table.concat(out)
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, following \input/\include
-- by default (see module docstring). Returns a 4th value: a list of
-- warning strings for any includes that couldn't be resolved.
-- @param opts table|nil { follow_includes = true, max_depth = 20 }
function M.count_file(path, config, opts)
  opts = opts or {}
  local follow = opts.follow_includes
  if follow == nil then
    follow = true
  end

  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()

  local warnings = {}
  local text = content
  if follow then
    local base_dir = path:match("^(.*)/[^/]+$") or "."
    local basename = path:match("([^/]+)$") or path
    local root_key = M.resolve_tex_path(basename, base_dir) or path
    local abs_seen = {}
    abs_seen[root_key] = true
    text = M.expand_includes(content, {
      base_dir = base_dir,
      seen = abs_seen,
      max_depth = opts.max_depth,
      warnings = warnings,
    })
  end

  local count, words, cleaned = M.count_text(text, config)
  return count, words, cleaned, warnings
end

--- Run the full pipeline over a Neovim buffer's contents, following
-- \input/\include relative to the buffer's own file by default. Returns
-- a 4th value: a list of warning strings for unresolved includes.
-- @param bufnr number|nil buffer handle (0 or nil = current buffer)
-- @param opts table|nil { follow_includes = true, max_depth = 20 }
function M.count_buffer(bufnr, config, opts)
  if vim == nil then
    error("count_latex_words: count_buffer() requires Neovim's 'vim' API")
  end
  opts = opts or {}
  local follow = opts.follow_includes
  if follow == nil then
    follow = true
  end

  bufnr = bufnr or 0
  local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
  local text = table.concat(lines, "\n")

  local warnings = {}
  if follow then
    local bufname = vim.api.nvim_buf_get_name(bufnr)
    local base_dir = (bufname ~= "" and bufname:match("^(.*)/[^/]+$")) or vim.fn.getcwd()
    local seen = {}
    if bufname ~= "" then
      local basename = bufname:match("([^/]+)$") or bufname
      local root_key = M.resolve_tex_path(basename, base_dir) or bufname
      seen[root_key] = true
    end
    text = M.expand_includes(text, {
      base_dir = base_dir,
      seen = seen,
      max_depth = opts.max_depth,
      warnings = warnings,
    })
  end

  local count, words, cleaned = M.count_text(text, config)
  return count, words, cleaned, warnings
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("MultiLatexWordCount", 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, _, _, warnings = M.count_buffer(0, config)
    vim.notify(string.format("Real words: %d", count))
    if warnings and #warnings > 0 then
      vim.notify("count_latex_words include warnings:\n" .. table.concat(warnings, "\n"), vim.log.levels.WARN)
    end
  end, { nargs = "?", complete = "file", desc = "Count real (visible) words in the current LaTeX buffer (Including other files)" })

  vim.api.nvim_create_user_command("SingleLatexWordCount", 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, _, _, warnings = M.count_buffer(0, config, { follow_includes = false })
    vim.notify(string.format("Real words: %d", count))
    if warnings and #warnings > 0 then
      vim.notify("count_latex_words include warnings:\n" .. table.concat(warnings, "\n"), vim.log.levels.WARN)
    end
  end, { nargs = "?", complete = "file", desc = "Count real (visible) words in the current LaTeX buffer" })
end
return M