const JSON_SIMPLE_ESCAPES = new Set(['"', '\\', '/', 'b', 'f', 'n', 'r', 't']);

/**
 * Detects a malformed-model collision where a TeX control sequence begins
 * with a character JSON would otherwise consume as a one-letter escape. This
 * intentionally avoids a finite command allow-list so custom macros and less
 * common controls (for example `\normalfont` and `\numberwithin`) survive.
 *
 * A capital letter immediately after the JSON escape is treated as ordinary
 * escaped prose (for example `\nHeading`) rather than TeX. One-letter TeX
 * controls followed by a braced argument are also protected.
 */
function latexCollision(source: string, slashIndex: number): boolean {
  const first = source[slashIndex + 1];
  if (!first || !JSON_SIMPLE_ESCAPES.has(first) || first === '"' || first === '\\' || first === '/') return false;
  const command = /^[A-Za-z@]+/u.exec(source.slice(slashIndex + 1))?.[0] ?? '';
  if (!command) return false;
  if (command.length === 1) return source[slashIndex + 2] === '{';
  const continuation = command[1];
  return continuation === '@' || (continuation !== undefined && continuation === continuation.toLowerCase());
}

function stripModelEnvelope(text: string): string {
  return text
    .replace(/\[LEARNITIVE_[A-Z0-9_]+\]/gu, '')
    .trim()
    .replace(/^```(?:json)?\s*/iu, '')
    .replace(/\s*```$/u, '')
    .trim();
}

/**
 * Makes model-authored JSON string contents safe for JSON.parse without
 * changing object structure or inventing semantic content. In particular,
 * LaTeX commands such as `\textbf` and `\ref` are protected from JSON's
 * one-letter escape sequences so they cannot silently turn into control chars.
 */
export function latexSafeAiJsonSource(text: string): string {
  let output = '';
  let inString = false;
  for (let index = 0; index < text.length; index += 1) {
    const char = text[index]!;
    if (!inString) {
      output += char;
      if (char === '"') inString = true;
      continue;
    }
    if (char === '"') {
      output += char;
      inString = false;
      continue;
    }
    if (char === '\n') { output += '\\n'; continue; }
    if (char === '\r') { output += '\\r'; continue; }
    if (char === '\t') { output += '\\t'; continue; }
    if (char !== '\\') {
      output += char;
      continue;
    }

    const next = text[index + 1];
    if (next === undefined) {
      output += '\\\\';
      continue;
    }
    if (next === 'u' && /^[0-9a-fA-F]{4}$/u.test(text.slice(index + 2, index + 6))) {
      output += text.slice(index, index + 6);
      index += 5;
      continue;
    }
    if (next === '"' || next === '\\' || next === '/') {
      output += `\\${next}`;
      index += 1;
      continue;
    }

    if (!JSON_SIMPLE_ESCAPES.has(next) || latexCollision(text, index)) {
      output += '\\\\';
      continue;
    }
    output += `\\${next}`;
    index += 1;
  }
  return output;
}

function rootStart(source: string, allowArray: boolean): number {
  const object = source.indexOf('{');
  const array = allowArray ? source.indexOf('[') : -1;
  if (object < 0) return array;
  if (array < 0) return object;
  return Math.min(object, array);
}

function balancedPrefix(source: string): { candidate: string; complete: boolean } | undefined {
  const first = source[0];
  if (first !== '{' && first !== '[') return undefined;
  const stack: Array<'}' | ']'> = [first === '{' ? '}' : ']'];
  let inString = false;
  let escaped = false;
  for (let index = 1; index < source.length; index += 1) {
    const char = source[index]!;
    if (inString) {
      if (escaped) escaped = false;
      else if (char === '\\') escaped = true;
      else if (char === '"') inString = false;
      continue;
    }
    if (char === '"') { inString = true; continue; }
    if (char === '{') stack.push('}');
    else if (char === '[') stack.push(']');
    else if (char === '}' || char === ']') {
      if (stack.at(-1) !== char) return undefined;
      stack.pop();
      if (stack.length === 0) return { candidate: source.slice(0, index + 1), complete: true };
    }
  }
  return { candidate: source, complete: false };
}

function repairTruncatedJson(source: string): string {
  const stack: Array<'}' | ']'> = [];
  let inString = false;
  let escaped = false;
  for (let index = 0; index < source.length; index += 1) {
    const char = source[index]!;
    if (inString) {
      if (escaped) escaped = false;
      else if (char === '\\') escaped = true;
      else if (char === '"') inString = false;
      continue;
    }
    if (char === '"') { inString = true; continue; }
    if (char === '{') stack.push('}');
    else if (char === '[') stack.push(']');
    else if ((char === '}' || char === ']') && stack.at(-1) === char) stack.pop();
  }
  let repaired = source.trimEnd();
  if (inString) repaired += '"';
  repaired = repaired.replace(/,\s*$/u, '');
  if (/:\s*$/u.test(repaired) || /[,[]\s*$/u.test(repaired)) repaired += 'null';
  repaired = repaired.replace(/,\s*$/u, '');
  repaired += [...stack].reverse().join('');
  return repaired.replace(/,\s*([}\]])/gu, '$1');
}

export function parseRecoverableAiJson(text: string, options: { allowArray?: boolean; recoverTruncation?: boolean } = {}): unknown {
  const cleaned = stripModelEnvelope(text);
  const start = rootStart(cleaned, options.allowArray === true);
  if (start < 0) throw new Error('No JSON object was returned');
  const source = cleaned.slice(start);
  const prefix = balancedPrefix(source);
  if (!prefix) throw new Error('Mismatched JSON delimiters');

  try {
    return JSON.parse(latexSafeAiJsonSource(prefix.candidate)) as unknown;
  } catch (firstError) {
    if (prefix.complete || options.recoverTruncation === false) throw firstError;
    return JSON.parse(latexSafeAiJsonSource(repairTruncatedJson(prefix.candidate))) as unknown;
  }
}

export function parseRecoverableAiJsonObject(text: string): Record<string, unknown> | undefined {
  try {
    const value = parseRecoverableAiJson(text, { recoverTruncation: true });
    return value && typeof value === 'object' && !Array.isArray(value) ? value as Record<string, unknown> : undefined;
  } catch {
    return undefined;
  }
}