import { ParserRuleContext } from "antlr4ng";
import { parseDocument } from "./parser.js";
import { QuixosCapabilityParser } from "./generated/QuixosCapabilityParser.js";

/** Offsets and columns use UTF-16, as do JavaScript and LSP. End is exclusive. */
export type SourceRange = { start: number; end: number };
export type SourceEdit = SourceRange & { text: string };
export type SyntaxNode = SourceRange & { kind: string; children: SyntaxNode[] };

export const parseQx = (source: string, fileName = "<memory>") => {
  const parsed = parseDocument(source, fileName);
  // ANTLR indexes Unicode code points; editors index UTF-16 code units.
  const offsets = [0];
  for (const character of source) offsets.push(offsets[offsets.length - 1]! + character.length);
  const offset = (index: number) => offsets[Math.max(0, index)] ?? source.length;
  const node = (context: ParserRuleContext): SyntaxNode => ({
    kind: QuixosCapabilityParser.ruleNames[context.ruleIndex]!,
    start: offset(context.start?.start ?? 0),
    end: offset((context.stop?.stop ?? -1) + 1),
    children: context.children.flatMap((child) => (child instanceof ParserRuleContext ? [node(child)] : [])),
  });
  return {
    source,
    fileName,
    root: node(parsed.tree),
    diagnostics: parsed.diagnostics.map((diagnostic) => {
      const line = source.split("\n")[diagnostic.line - 1] ?? "";
      return { ...diagnostic, column: [...line].slice(0, diagnostic.column).join("").length };
    }),
    tokens: parsed.tokens
      .getTokens()
      .filter((token) => token.type !== -1)
      .map((token) => ({
        kind: QuixosCapabilityParser.symbolicNames[token.type] ?? "token",
        start: offset(token.start),
        end: offset(token.stop + 1),
        text: token.text ?? "",
        trivia: token.channel !== 0,
      })),
  };
};

/** Edits refer to one immutable source snapshot; overlapping edits are errors. */
export const applySourceEdits = (source: string, edits: readonly SourceEdit[]) => {
  const sorted = [...edits].sort((a, b) => a.start - b.start || a.end - b.end);
  let end = 0;
  let previousStart = -1;
  let result = "";
  for (const edit of sorted) {
    if (
      !Number.isInteger(edit.start) ||
      !Number.isInteger(edit.end) ||
      edit.start < end ||
      edit.start === previousStart ||
      edit.end < edit.start ||
      edit.end > source.length
    ) {
      throw new Error("Invalid or overlapping source edits");
    }
    result += source.slice(end, edit.start) + edit.text;
    end = edit.end;
    previousStart = edit.start;
  }
  return result + source.slice(end);
};

export const walkSyntax = function* (node: SyntaxNode): Generator<SyntaxNode> {
  yield node;
  for (const child of node.children) yield* walkSyntax(child);
};

export const lintQx = (source: string, fileName = "<memory>") => {
  const syntax = parseQx(source, fileName);
  const diagnostics = syntax.diagnostics.map((entry) => ({ ...entry, severity: "error" as "error" | "warning" }));
  if (diagnostics.length) return diagnostics;
  const seen = new Set<string>();
  for (const node of walkSyntax(syntax.root)) {
    if (node.kind !== "sourceImportDecl") continue;
    const literal = node.children.find((child) => child.kind === "stringLiteral")!;
    const importPath: string = JSON.parse(source.slice(literal.start, literal.end));
    let message: string | undefined;
    let code = "invalid-source-import";
    try {
      validateQxImportPath(importPath);
    } catch (error) {
      message = (error as Error).message;
    }
    if (!message && seen.has(importPath)) {
      code = "duplicate-source-import";
      message = `Repeated local import ${importPath}`;
    }
    seen.add(importPath);
    if (message) {
      const prefix = source.slice(0, node.start);
      diagnostics.push({
        phase: "syntax",
        code,
        message,
        fileName,
        line: prefix.split("\n").length,
        column: prefix.length - prefix.lastIndexOf("\n") - 1,
        severity: code === "duplicate-source-import" ? "warning" : "error",
      });
    }
  }
  return diagnostics;
};

/** Conservative formatter: indentation only, preserving strings and comments verbatim. */
export const formatQx = (source: string) => {
  const syntax = parseQx(source);
  if (syntax.diagnostics.length) throw new Error("Cannot format QX with syntax errors");
  const edits: SourceEdit[] = [];
  let depth = 0;
  let lineStart = 0;
  for (const token of syntax.tokens) {
    if (token.kind === "WS") continue;
    lineStart = source.lastIndexOf("\n", token.start - 1) + 1;
    if (/^[ \t]*$/.test(source.slice(lineStart, token.start))) {
      const indentation = "  ".repeat(Math.max(0, depth - (token.kind === "RBRACE" ? 1 : 0)));
      if (source.slice(lineStart, token.start) !== indentation)
        edits.push({ start: lineStart, end: token.start, text: indentation });
    }
    if (!token.trivia) {
      if (token.kind === "LBRACE") depth++;
      if (token.kind === "RBRACE") depth--;
    }
  }
  return applySourceEdits(source, edits);
};

export const addWorkspaceImport = (source: string, importPath: string) => {
  validateQxImportPath(importPath);
  const syntax = parseQx(source);
  if (syntax.diagnostics.length || syntax.root.children[0]?.kind !== "workspaceDecl")
    throw new Error("Expected a syntactically valid workspace");
  for (const node of walkSyntax(syntax.root)) {
    if (node.kind === "sourceImportDecl") {
      const literal = node.children.find((child) => child.kind === "stringLiteral")!;
      if (JSON.parse(source.slice(literal.start, literal.end)) === importPath) return source;
    }
  }
  const brace = syntax.tokens.find((token) => token.kind === "LBRACE")!;
  const newline = source.includes("\r\n") ? "\r\n" : "\n";
  return applySourceEdits(source, [
    { start: brace.end, end: brace.end, text: `${newline}  import ${JSON.stringify(importPath)};` },
  ]);
};

export const validateQxImportPath = (value: string) => {
  if (
    !/^(?:[A-Za-z0-9_-][A-Za-z0-9_.-]*\/)*[A-Za-z0-9_-][A-Za-z0-9_.-]*\.qx$/.test(value) ||
    value.split("/").some((part) => part === "." || part === "..")
  )
    throw new Error(`Invalid repository-relative QX import path: ${value}`);
};