/**
 * Cypher rewrite pass that injects provenance stamps into operator-supplied
 * write_neo4j_cypher. Sibling to cypher-validate.ts: validate first,
 * rewrite second, run third.
 *
 * Doctrine context: wrapped writers stamp every CREATE node
 * structurally via writeNodeWithEdges; raw cypher relied on
 * prompt-internalised discipline. This module closes the asymmetry — the
 * rewriter appends `SET <alias>.createdAt = $__autoStartTimestamp,
 * <alias>.createdByAgent = $__autoAgent, <alias>.createdByTool =
 * 'graph-cypher-write', <alias>.createdBySession = $__autoSession` to every
 * CREATE/MERGE clause that introduces a new labeled node. Caller supplies
 * the three `$__auto*` params at tx.run time.
 *
 * MERGE handling preserves the create-vs-match distinction. CREATE always
 * stamps (every CREATE introduces a new node by definition). MERGE injects
 * `ON CREATE SET` so an existing-node match path keeps the original writer's
 * provenance intact — overwriting it would destroy forensic attribution.
 *
 * Node-only stamping. Edge aliases (`[r:R]`) are not auto-stamped — matches
 * the wrapped writer's writeNodeWithEdges, which stamps the node only.
 *
 * Idempotent: re-applying the rewriter to its own output is a no-op. The
 * marker that proves "already stamped" is `<alias>.createdBySession =
 * $__autoSession` — the unique combination of property name + auto param.
 *
 * String literals and comments are redacted (chars replaced with spaces,
 * offsets preserved) so a `CREATE (fake:Foo)` inside a comment or quoted
 * string cannot mistakenly trip the rewriter.
 */

export interface RewriteResult {
  cypher: string;
  /** Number of node aliases that received a fresh stamp injection. */
  stampsAppended: number;
}

const AUTO_PARAM_AGENT = "$__autoAgent";
const AUTO_PARAM_SESSION = "$__autoSession";
const AUTO_PARAM_TIMESTAMP = "$__autoStartTimestamp";
const AUTO_TOOL_LITERAL = "'graph-cypher-write'";

const KEYWORD_TERMINATORS = new Set([
  "SET",
  "ON",
  "MATCH",
  "MERGE",
  "CREATE",
  "RETURN",
  "WITH",
  "WHERE",
  "REMOVE",
  "DELETE",
  "DETACH",
  "UNION",
  "ORDER",
  "LIMIT",
  "SKIP",
  "FOREACH",
  "UNWIND",
  "CALL",
  "USE",
  "LOAD",
  "YIELD",
]);

function buildRedacted(cypher: string): string {
  let out = "";
  let i = 0;
  while (i < cypher.length) {
    const c = cypher[i];
    if (c === "/" && cypher[i + 1] === "/") {
      while (i < cypher.length && cypher[i] !== "\n") {
        out += " ";
        i++;
      }
      continue;
    }
    if (c === "/" && cypher[i + 1] === "*") {
      out += "  ";
      i += 2;
      while (i < cypher.length && !(cypher[i] === "*" && cypher[i + 1] === "/")) {
        out += " ";
        i++;
      }
      if (i < cypher.length) {
        out += "  ";
        i += 2;
      }
      continue;
    }
    if (c === "'" || c === '"') {
      const quote = c;
      out += c;
      i++;
      while (i < cypher.length && cypher[i] !== quote) {
        if (cypher[i] === "\\" && i + 1 < cypher.length) {
          out += "  ";
          i += 2;
          continue;
        }
        out += " ";
        i++;
      }
      if (i < cypher.length) {
        out += cypher[i];
        i++;
      }
      continue;
    }
    out += c;
    i++;
  }
  return out;
}

function findKeywordHits(
  redacted: string,
): Array<{ kw: "CREATE" | "MERGE"; afterKw: number }> {
  const re = /\b(CREATE|MERGE)\b/gi;
  const results: Array<{ kw: "CREATE" | "MERGE"; afterKw: number }> = [];
  for (const m of redacted.matchAll(re)) {
    if (m.index === undefined) continue;
    const upper = m[1].toUpperCase();
    // Skip `CREATE` that's part of `ON CREATE SET` (a MERGE sub-clause, not
    // an introducing CREATE). Without this, the auto-stamp's own injected
    // `ON CREATE SET` would register as a separate CREATE hit on the next
    // rewriter pass, breaking idempotency by truncating the stamp window.
    if (upper === "CREATE") {
      const prefix = redacted.slice(Math.max(0, m.index - 8), m.index);
      if (/\bON\s+$/i.test(prefix)) continue;
    }
    results.push({
      kw: upper as "CREATE" | "MERGE",
      afterKw: m.index + m[0].length,
    });
  }
  return results;
}

const TERMINATOR_PROBE = /^(\s*)(?:;|([A-Za-z]+))/;

/**
 * Returns true iff position `i` in `redacted` is at a word boundary — i.e.,
 * the immediately preceding character is not part of an identifier (letters,
 * digits, underscore). Without this guard, TERMINATOR_PROBE matches the
 * trailing letters of a multi-char identifier as a keyword (e.g. `Session`
 * ending in `on` would match the `ON` keyword and truncate the scan
 * mid-identifier).
 */
function atWordBoundary(redacted: string, i: number): boolean {
  if (i === 0) return true;
  return !/[A-Za-z0-9_]/.test(redacted[i - 1]);
}

/**
 * Returns the index just past the closing `)` of the CREATE/MERGE pattern,
 * or null if no node pattern follows the keyword (e.g. CREATE INDEX, CREATE
 * CONSTRAINT — neither would have reached the rewriter, but defensive).
 */
function findPatternEnd(redacted: string, afterKw: number): number | null {
  let i = afterKw;
  while (i < redacted.length && /\s/.test(redacted[i])) i++;
  if (i >= redacted.length || redacted[i] !== "(") return null;
  let parenDepth = 0;
  let bracketDepth = 0;
  while (i < redacted.length) {
    if (parenDepth === 0 && bracketDepth === 0 && i > afterKw) {
      const tmatch = redacted.slice(i).match(TERMINATOR_PROBE);
      if (tmatch) {
        const wsLen = tmatch[1].length;
        const ch = redacted[i + wsLen];
        if (ch === ";") return i + wsLen;
        const word = tmatch[2];
        if (
          word &&
          KEYWORD_TERMINATORS.has(word.toUpperCase()) &&
          atWordBoundary(redacted, i + wsLen)
        ) {
          return i + wsLen;
        }
      }
    }
    const c = redacted[i];
    if (c === "(") parenDepth++;
    else if (c === ")") parenDepth--;
    else if (c === "[") bracketDepth++;
    else if (c === "]") bracketDepth--;
    i++;
  }
  return i;
}

function extractLabeledAliases(redactedSlice: string): string[] {
  const out: string[] = [];
  const seen = new Set<string>();
  const re = /\(\s*([A-Za-z_][A-Za-z0-9_]*)\s*:\s*[A-Z]/g;
  for (const m of redactedSlice.matchAll(re)) {
    const alias = m[1];
    if (!seen.has(alias)) {
      seen.add(alias);
      out.push(alias);
    }
  }
  return out;
}

function aliasAlreadyStamped(slice: string, alias: string): boolean {
  // Per-pattern scope (review F1). The `slice` covers this CREATE/MERGE
  // pattern's possible SET / ON CREATE SET clauses up to the next CREATE/
  // MERGE boundary — scoping the marker check globally would mis-classify
  // a re-introduction of the same alias name across two distinct patterns
  // as "already stamped" (the first pattern's stamp would shadow the second).
  const escaped = alias.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
  const marker = new RegExp(
    `\\b${escaped}\\.createdBySession\\s*=\\s*\\$__autoSession\\b`,
  );
  return marker.test(slice);
}

function buildAssignments(aliases: string[]): string {
  return aliases
    .flatMap((a) => [
      `${a}.createdAt = ${AUTO_PARAM_TIMESTAMP}`,
      `${a}.createdByAgent = ${AUTO_PARAM_AGENT}`,
      `${a}.createdByTool = ${AUTO_TOOL_LITERAL}`,
      `${a}.createdBySession = ${AUTO_PARAM_SESSION}`,
    ])
    .join(", ");
}

const ON_CREATE_SET_PROBE = /^ON\s+CREATE\s+SET\s+/i;
const ON_MATCH_SET_PROBE = /^ON\s+MATCH\s+SET\b/i;

function findOnCreateSetEnd(
  redacted: string,
  fromIdx: number,
): { end: number } | null {
  let i = fromIdx;
  while (i < redacted.length && /\s/.test(redacted[i])) i++;
  const m = redacted.slice(i).match(ON_CREATE_SET_PROBE);
  if (!m) return null;
  let j = i + m[0].length;
  let parenDepth = 0;
  let bracketDepth = 0;
  while (j < redacted.length) {
    if (parenDepth === 0 && bracketDepth === 0) {
      const tmatch = redacted.slice(j).match(TERMINATOR_PROBE);
      if (tmatch) {
        const wsLen = tmatch[1].length;
        const ch = redacted[j + wsLen];
        if (ch === ";") return { end: j + wsLen };
        const word = tmatch[2];
        if (
          word &&
          KEYWORD_TERMINATORS.has(word.toUpperCase()) &&
          atWordBoundary(redacted, j + wsLen)
        ) {
          return { end: j + wsLen };
        }
      }
    }
    const c = redacted[j];
    if (c === "(") parenDepth++;
    else if (c === ")") parenDepth--;
    else if (c === "[") bracketDepth++;
    else if (c === "]") bracketDepth--;
    j++;
  }
  return { end: j };
}

function findOnMatchSetStart(redacted: string, fromIdx: number): number | null {
  let i = fromIdx;
  while (i < redacted.length && /\s/.test(redacted[i])) i++;
  if (ON_MATCH_SET_PROBE.test(redacted.slice(i))) {
    return i;
  }
  return null;
}

export function rewriteWithProvenanceStamps(cypher: string): RewriteResult {
  const redacted = buildRedacted(cypher);
  const hits = findKeywordHits(redacted);

  type Edit = { at: number; insert: string };
  const edits: Edit[] = [];
  let stampsAppended = 0;

  for (let h = 0; h < hits.length; h++) {
    const { kw, afterKw } = hits[h];
    const patternEnd = findPatternEnd(redacted, afterKw);
    if (patternEnd === null) continue;

    let pStart = afterKw;
    while (pStart < redacted.length && /\s/.test(redacted[pStart])) pStart++;

    const patternSlice = redacted.slice(pStart, patternEnd);
    const labeled = extractLabeledAliases(patternSlice);
    // Per-pattern stamp window: from this pattern's start through the next
    // CREATE/MERGE boundary (or end of cypher). Scoping the "already stamped"
    // check globally would mis-classify a re-introduction of the same alias
    // name across two distinct patterns (review F1).
    const nextHitAfterKw =
      h + 1 < hits.length ? hits[h + 1].afterKw : redacted.length;
    const stampWindow = redacted.slice(pStart, nextHitAfterKw);
    const aliases = labeled.filter((a) => !aliasAlreadyStamped(stampWindow, a));
    if (aliases.length === 0) continue;

    if (kw === "CREATE") {
      // Trailing space matters: when patternEnd lands at a clause keyword
      // (e.g. ` WITH` immediately after `(n:Person)`), our `findPatternEnd`
      // returns the keyword's position. Without a trailing space, the
      // injected `$__autoSession` runs into the keyword and Cypher parses
      // `$__autoSessionWITH` as one parameter name.
      edits.push({
        at: patternEnd,
        insert: ` SET ${buildAssignments(aliases)} `,
      });
      stampsAppended += aliases.length;
      continue;
    }

    const existingOnCreate = findOnCreateSetEnd(redacted, patternEnd);
    if (existingOnCreate) {
      edits.push({
        at: existingOnCreate.end,
        insert: `, ${buildAssignments(aliases)} `,
      });
      stampsAppended += aliases.length;
      continue;
    }
    const onMatchStart = findOnMatchSetStart(redacted, patternEnd);
    const insertAt = onMatchStart ?? patternEnd;
    edits.push({
      at: insertAt,
      insert: ` ON CREATE SET ${buildAssignments(aliases)} `,
    });
    stampsAppended += aliases.length;
  }

  edits.sort((a, b) => b.at - a.at);
  let result = cypher;
  for (const edit of edits) {
    result = result.slice(0, edit.at) + edit.insert + result.slice(edit.at);
  }

  return { cypher: result, stampsAppended };
}