/**
 * llmindex.io — answer/confidence extraction and grading (robust cascade)
 * Author:  Simon-Pierre Boucher
 * Contact: contact@spboucher.ai
 * License: Proprietary — © Simon-Pierre Boucher, all rights reserved
 *
 * Extraction must never confound formatting with ability: models wrap answers
 * in markdown bold, backticks, LaTeX \boxed{}, or drop the tag entirely.
 * A lenient, ordered cascade extracts the intended answer; grading then
 * normalizes aggressively. Format compliance is measured by its own domain
 * (instruction_following), not silently through every other domain.
 */

/** Instruction appended to every graded item prompt. */
export const ANSWER_FORMAT_INSTRUCTIONS =
  'End your reply with exactly two plain-text lines (no markdown, no extra text after them):\n' +
  'ANSWER: <your final answer only>\n' +
  'CONFIDENCE: <integer 0-100, how confident you are that your answer is correct>';

/** Instruction for items whose answer is a JSON payload or multi-line output. */
export const BLOCK_ANSWER_FORMAT_INSTRUCTIONS =
  'Give your final answer inside ONE fenced code block (```), containing exactly the required ' +
  'content and nothing else. After the code block, add one plain-text line:\n' +
  'CONFIDENCE: <integer 0-100, how confident you are that your answer is correct>';

export interface ExtractedAnswer {
  answer: string | null;
  /** Confidence in [0,1], null if the model did not report one. */
  confidence: number | null;
}

/** Strip markdown decorations that models wrap around tags and values. */
function stripMd(s: string): string {
  return s
    .replace(/\*\*|__|~~|`+/g, '')
    .replace(/^\s*[#>*-]+\s*/, '')
    .trim();
}

/** Unwrap LaTeX decorations: \boxed{x}, \text{x}, $x$, \( x \). */
function stripLatex(s: string): string {
  let out = s.trim();
  const boxed = out.match(/\\boxed\s*\{([^{}]*)\}/);
  if (boxed?.[1] != null) out = boxed[1];
  out = out
    .replace(/\\text\s*\{([^{}]*)\}/g, '$1')
    .replace(/\\mathrm\s*\{([^{}]*)\}/g, '$1')
    .replace(/^\$+|\$+$/g, '')
    .replace(/^\\\(|\\\)$/g, '')
    .trim();
  return out;
}

function parseConfidence(text: string): number | null {
  const lines = text.split('\n').map(stripMd);
  const matches = lines
    .map((l) => l.match(/^CONFIDENCE\s*[:=]?\s*(\d{1,3})\s*%?\s*\.?$/i))
    .filter((m): m is RegExpMatchArray => m !== null);
  const last = matches[matches.length - 1];
  if (!last) return null;
  const v = Number(last[1]);
  return Number.isFinite(v) && v >= 0 && v <= 100 ? v / 100 : null;
}

/**
 * Cascade for line answers:
 *  1. last markdown-stripped line matching "ANSWER: x" (also FINAL ANSWER / Answer =)
 *  2. \boxed{x} anywhere (last occurrence)
 *  3. null — graders may apply their own last-resort fallback (numeric only)
 */
export function extractAnswer(text: string): ExtractedAnswer {
  const confidence = parseConfidence(text);
  const lines = text.split('\n').map(stripMd);
  const tagMatches = lines
    .map((l) => l.match(/^(?:FINAL\s+)?ANSWER\s*[:=]\s*(.+?)\s*$/i))
    .filter((m): m is RegExpMatchArray => m !== null);
  const lastTag = tagMatches[tagMatches.length - 1];
  if (lastTag?.[1]) {
    return { answer: stripLatex(lastTag[1]), confidence };
  }
  const boxed = [...text.matchAll(/\\boxed\s*\{([^{}]*)\}/g)];
  const lastBoxed = boxed[boxed.length - 1];
  if (lastBoxed?.[1]) {
    return { answer: lastBoxed[1].trim(), confidence };
  }
  return { answer: null, confidence };
}

/**
 * Extraction for block answers (JSON call sequences, multi-line terminal
 * output): last fenced code block; falls back to text after the last
 * "ANSWER:" tag when the model skipped the fence.
 */
export function extractBlockAnswer(text: string): ExtractedAnswer {
  const confidence = parseConfidence(text);
  const fences = [...text.matchAll(/```[a-zA-Z]*\r?\n([\s\S]*?)```/g)];
  const last = fences[fences.length - 1];
  if (last?.[1] != null && last[1].trim().length > 0) {
    return { answer: last[1].replace(/\s+$/, ''), confidence };
  }
  const tag = text.match(/(?:FINAL\s+)?ANSWER\s*[:=]\s*([\s\S]+)$/i);
  if (tag?.[1]) {
    return { answer: tag[1].replace(/CONFIDENCE\s*[:=][\s\S]*$/i, '').trim(), confidence };
  }
  return { answer: null, confidence };
}

export function normalizeAnswer(raw: string): string {
  return stripLatex(stripMd(raw))
    .toLowerCase()
    .replace(/[""''`´]/g, '')
    .replace(/[\s.,;:!?]+$/g, '')
    .replace(/-/g, ' ') // hyphen/space orthography variants (vingt-deux ≡ vingt deux)
    .replace(/\s+/g, ' ')
    .trim();
}

function parseNumeric(raw: string): number | null {
  let cleaned = stripLatex(stripMd(raw))
    .replace(/[$€£%]/g, '')
    .replace(/[\u00a0\u202f]/g, ' ')
    .trim();
  // thousands separators: "1,234,567" or "1 234 567"
  cleaned = cleaned.replace(/(\d)[ ,](?=\d{3}(\D|$))/g, '$1');
  cleaned = cleaned.replace(/[a-zA-Z]+$/g, '').trim(); // trailing units
  const m = cleaned.match(/-?\d+(?:\.\d+)?/);
  if (!m) return null;
  if (cleaned !== m[0]) {
    // reject if extra numeric garbage surrounds ("12 or 13")
    const rest = cleaned.replace(m[0], '');
    if (/\d/.test(rest)) return null;
  }
  const v = Number(m[0]);
  return Number.isFinite(v) ? v : null;
}

/** Canonicalize JSON: stable key order, numbers as-is, whitespace-free. */
export function canonicalJson(value: unknown): string {
  if (Array.isArray(value)) return `[${value.map(canonicalJson).join(',')}]`;
  if (value !== null && typeof value === 'object') {
    const entries = Object.entries(value as Record<string, unknown>)
      .filter(([, v]) => v !== undefined)
      .sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0));
    return `{${entries.map(([k, v]) => `${JSON.stringify(k)}:${canonicalJson(v)}`).join(',')}}`;
  }
  return JSON.stringify(value);
}

function tryParseJson(raw: string): unknown | undefined {
  const cleaned = raw
    .replace(/```[a-zA-Z]*\r?\n?/g, '')
    .replace(/```/g, '')
    .trim();
  const start = cleaned.search(/[[{]/);
  if (start === -1) return undefined;
  const candidate = cleaned.slice(start);
  try {
    return JSON.parse(candidate);
  } catch {
    // trim trailing prose after the JSON payload
    for (let end = candidate.length; end > 1; end--) {
      const c = candidate[end - 1];
      if (c === ']' || c === '}') {
        try {
          return JSON.parse(candidate.slice(0, end));
        } catch {
          /* keep scanning */
        }
      }
    }
    return undefined;
  }
}

export type GradingMode = 'exact' | 'numeric' | 'json' | 'lines' | 'constraints';

/**
 * Constraint-stack spec (instruction_following): the answer key is a JSON
 * spec; ANY output satisfying every constraint is correct. All checks are
 * mechanical — no judges.
 */
export interface ConstraintSpec {
  wordCount?: number;
  startsWithWord?: string;
  endsWithWord?: string;
  includeWordExactly?: Array<{ word: string; count: number }>;
  forbiddenLetter?: string;
  allLowercase?: boolean;
}

export function checkConstraints(output: string, spec: ConstraintSpec): boolean {
  const text = output.trim();
  const words = text
    .toLowerCase()
    .replace(/[.,;:!?"()]/g, ' ')
    .split(/\s+/)
    .filter(Boolean);
  if (spec.wordCount !== undefined && words.length !== spec.wordCount) return false;
  if (spec.startsWithWord && words[0] !== spec.startsWithWord.toLowerCase()) return false;
  if (spec.endsWithWord && words[words.length - 1] !== spec.endsWithWord.toLowerCase()) return false;
  if (spec.includeWordExactly) {
    for (const { word, count } of spec.includeWordExactly) {
      if (words.filter((w) => w === word.toLowerCase()).length !== count) return false;
    }
  }
  if (spec.forbiddenLetter && text.toLowerCase().includes(spec.forbiddenLetter.toLowerCase()))
    return false;
  if (spec.allLowercase && text !== text.toLowerCase()) return false;
  return true;
}

/** Grade an extracted answer against a key. */
export function gradeAnswer(extracted: string | null, answerKey: string, grading: GradingMode): boolean {
  if (extracted === null) return false;
  switch (grading) {
    case 'numeric': {
      const got = parseNumeric(extracted);
      const want = parseNumeric(answerKey);
      if (got === null || want === null) return false;
      return Math.abs(got - want) <= Math.max(1e-9, Math.abs(want) * 1e-6);
    }
    case 'json': {
      const got = tryParseJson(extracted);
      const want = tryParseJson(answerKey);
      if (got === undefined || want === undefined) return false;
      return canonicalJson(got) === canonicalJson(want);
    }
    case 'constraints': {
      try {
        return checkConstraints(extracted, JSON.parse(answerKey) as ConstraintSpec);
      } catch {
        return false;
      }
    }
    case 'lines': {
      const clean = (s: string): string =>
        s
          .split('\n')
          .map((l) => l.replace(/\s+$/g, ''))
          .filter((l, i, arr) => !(l === '' && (i === 0 || i === arr.length - 1)))
          .join('\n')
          .trim();
      return clean(extracted) === clean(answerKey);
    }
    default:
      return normalizeAnswer(extracted) === normalizeAnswer(answerKey);
  }
}