// After each file edit, list the comment lines it added and ask the model to
// justify or remove them. Everything runs locally: no binary, no network.
import { readFileSync } from "node:fs"
import { extname, isAbsolute, relative, resolve } from "node:path"

const EDIT_TOOLS = new Set(["edit", "write", "apply_patch"])

const LINE_MARKERS: Record<string, string[]> = {
  "//": [".rs", ".c", ".h", ".cc", ".cpp", ".cxx", ".hpp", ".hh", ".js", ".jsx", ".mjs", ".cjs",
         ".ts", ".tsx", ".mts", ".cts", ".go", ".java", ".kt", ".swift", ".cs", ".scala", ".dart",
         ".zig", ".proto", ".fbs", ".glsl", ".wgsl"],
  "#": [".py", ".pyi", ".sh", ".bash", ".zsh", ".rb", ".pl", ".r", ".jl", ".nim", ".cmake"],
  "--": [".lua", ".sql", ".hs"],
  ";": [".el", ".clj", ".lisp", ".scm"],
}

const MARKER_BY_EXT = new Map<string, string>()
for (const [marker, exts] of Object.entries(LINE_MARKERS)) {
  for (const ext of exts) MARKER_BY_EXT.set(ext, marker)
}

// Comments that tools or the user's own rules require.
const EXEMPT = [
  /^#!/,
  /\bTODO:/,
  /\bSAFETY:/,
  /SPDX-License-Identifier/,
  /-\*-.*-\*-/,
  /\b(noqa|type:\s*ignore|pyright:|mypy:|pylint:|ruff:|fmt:\s*(on|off|skip))/,
  /\b(eslint|prettier|biome|tslint)-/,
  /@ts-(ignore|expect-error|nocheck)/,
  /\b(NOLINT|clang-format|shellcheck)\b/,
  /^\/\/\s*(go:|\+build)/,
  /^\/\/\/\s*</,
]

function markerFor(path: string): string | undefined {
  const name = path.split("/").pop() ?? ""
  if (name === "CMakeLists.txt") return "#"
  return MARKER_BY_EXT.get(extname(path).toLowerCase())
}

// Index where a comment starts on this line, ignoring markers inside quotes.
function commentStart(line: string, marker: string): number {
  let quote = ""
  for (let i = 0; i < line.length; i++) {
    const ch = line[i]
    if (quote) {
      if (ch === "\\") i++
      else if (ch === quote) quote = ""
      continue
    }
    if (ch === "'" && marker === "//") {
      // Only a char literal opens a quote here, so a Rust lifetime does not.
      const close = line.indexOf("'", i + 1)
      if (close > i + 1 && close <= i + 3) i = close
      continue
    }
    if (ch === '"' || ch === "'" || ch === "`") {
      quote = ch
      continue
    }
    if (!line.startsWith(marker, i)) continue
    if (marker === "#" && i > 0 && !/\s/.test(line[i - 1])) continue
    if (marker === "//" && i > 0 && line[i - 1] === ":") continue
    return i
  }
  return -1
}

function commentText(line: string, marker: string): string | undefined {
  const trimmed = line.trim()
  if (marker === "//" && (trimmed.startsWith("/*") || (/^\*(\s|\/|$)/.test(trimmed) && !trimmed.endsWith(";")))) {
    return trimmed.replace(/^[/*\s]+$/, "") || undefined
  }
  if (marker === "#" && /^("""|''')/.test(trimmed)) return trimmed
  const at = commentStart(line, marker)
  if (at < 0) return undefined
  const text = line.slice(at).trim()
  if (text.replace(/^[^\w]+/, "") === "") return undefined
  return text
}

function readOrEmpty(path: string): string {
  try {
    return readFileSync(path, "utf8")
  } catch {
    return ""
  }
}

function targetPaths(tool: string, args: any, base: string): string[] {
  const raw: string[] = []
  if (tool === "apply_patch") {
    const text: string = args?.patchText ?? args?.patch ?? ""
    for (const m of text.matchAll(/^\*\*\* (?:Add|Update) File: (.+)$/gm)) raw.push(m[1].trim())
    for (const m of text.matchAll(/^\*\*\* Move to: (.+)$/gm)) raw.push(m[1].trim())
  } else if (args?.filePath) {
    raw.push(args.filePath)
  }
  return raw.map((p) => (isAbsolute(p) ? p : resolve(base, p)))
}

function addedComments(before: string, after: string, marker: string): { line: number; text: string }[] {
  const remaining = new Map<string, number>()
  for (const l of before.split("\n")) remaining.set(l, (remaining.get(l) ?? 0) + 1)
  const found: { line: number; text: string }[] = []
  after.split("\n").forEach((l, i) => {
    const left = remaining.get(l) ?? 0
    if (left > 0) {
      remaining.set(l, left - 1)
      return
    }
    const text = commentText(l, marker)
    if (text && !EXEMPT.some((re) => re.test(text))) found.push({ line: i + 1, text })
  })
  return found
}

const RULES = `Check each one against the comment rules:
- Keep a comment only where the why is not obvious from the code.
- Remove comments that restate what the code does, narrate the change, describe previous behavior, or explain a departure from an earlier approach. Say those things in the chat instead.
- Prefer one short line to a block. No semicolons or em-dashes in prose comments.
- Deferred work keeps a literal TODO: marker.
Remove the ones that fail with another edit now. If a flagged comment is warranted, leave it and do not mention this check.`

export const CommentCheck = async ({ directory }: { directory: string }) => {
  const pending = new Map<string, Map<string, string>>()

  return {
    "tool.execute.before": async (input: { tool: string; callID: string }, output: { args: any }) => {
      if (!EDIT_TOOLS.has(input.tool)) return
      const snapshot = new Map<string, string>()
      for (const p of targetPaths(input.tool, output.args, directory)) {
        if (markerFor(p)) snapshot.set(p, readOrEmpty(p))
      }
      if (snapshot.size) pending.set(input.callID, snapshot)
    },

    "tool.execute.after": async (
      input: { tool: string; callID: string; args: any },
      output: { output: string },
    ) => {
      const snapshot = pending.get(input.callID)
      if (!snapshot) return
      pending.delete(input.callID)

      const report: string[] = []
      for (const [path, before] of snapshot) {
        const marker = markerFor(path)!
        for (const c of addedComments(before, readOrEmpty(path), marker)) {
          report.push(`  ${relative(directory, path)}:${c.line}  ${c.text}`)
        }
      }
      if (!report.length) return

      output.output += `\n\n<comment-check>\nThis edit added ${report.length} comment line(s):\n${report.join("\n")}\n\n${RULES}\n</comment-check>`
    },
  }
}