/**
 * Turning model text into one action, and engine results into text.
 *
 * A port of `envs/alert_triage/alert_triage/protocol.py`. The one rule, as in
 * wordle: **never throw**. A malformed reply is a thing the model did, not an
 * error in the harness. It becomes a rejection that costs a turn.
 *
 * The parse rules are pinned because this is where two runtimes disagree, and
 * CI hashes a committed corpus of reply strings through this parser and the
 * engine on both sides (digest `protocol` in CONFORMANCE.txt):
 *
 *   1. If a ```json fence appears anywhere, its contents are the candidate.
 *      A bare `{` before the fence does not win.
 *   2. Otherwise the candidate is the first string-aware balanced `{…}` span:
 *      braces inside JSON strings do not count, backslash escapes are honoured.
 *   3. If the candidate is not valid JSON, or is JSON but not an object, the
 *      reply is rejected. There is no further scanning for a second candidate.
 *   4. Enum values are matched exact-case. Cites are deduplicated in order and
 *      matched exact-case against what the engine has shown.
 *
 * Why this file carries its own JSON reader rather than `JSON.parse`: the
 * engine's rejection reasons quote the offending value with Python's `repr`
 * (`unknown action None`, `month 3.0 is not in the twelve-month window`), and
 * those strings are hashed. `JSON.parse` cannot tell `3` from `3.0`, so the
 * reader keeps each number's source text, and `pyRepr` formats it the way
 * Python would. The reader accepts exactly what `json.loads` accepts — with
 * NaN/Infinity refused, as the Python does via `parse_constant`.
 */

import { MAX_TURNS } from './generator';
import type { Counterparty, Document, PriorAlert, Screen, Txn } from './generator';
import { DOC_MINUTES, LOOKUP_MINUTES, SCREEN_MINUTES, TURN_MINUTES } from './score';

export const ACTIONS = ['lookup', 'close', 'escalate'] as const;
export const LOOKUPS = ['history', 'counterparty', 'prior_alerts', 'documents'] as const;
export const CLOSE_REASONS = ['CONSISTENT_WITH_PROFILE', 'DOCUMENTED_SOURCE_OF_FUNDS', 'PREVIOUSLY_REVIEWED', 'RULE_ARTEFACT'] as const;
export const TYPOLOGIES = ['STRUCTURING', 'FUNNEL_ACCOUNT', 'MONEY_MULE', 'RAPID_MOVEMENT', 'UNKNOWN'] as const;

export type Action = (typeof ACTIONS)[number];
export type Lookup = (typeof LOOKUPS)[number];
export type CloseReason = (typeof CLOSE_REASONS)[number];
export type Typology = (typeof TYPOLOGIES)[number];

/* ----------------------------------------------------------- JSON values */

/**
 * A number as `json.loads` would hold it: `int` when the token had no
 * fraction or exponent, `float` otherwise. `raw` is the token's own text.
 */
export class JsonNumber {
  constructor(
    readonly raw: string,
    readonly value: number,
  ) {}
  get isInt(): boolean {
    return !/[.eE]/.test(this.raw);
  }
}

/** Objects are Maps so a key named `__proto__` is data, as it is in Python. */
export type JsonValue = null | boolean | string | JsonNumber | JsonValue[] | Map<string, JsonValue>;

export const isJsonObject = (v: JsonValue | undefined): v is Map<string, JsonValue> => v instanceof Map;

class JsonError extends Error {}

/**
 * `json.loads`, to the letter that matters: RFC 8259 grammar, `NaN` and
 * `Infinity` refused, control characters refused inside strings, lone
 * surrogate escapes allowed, trailing content refused. Throws `JsonError`.
 */
export function pyJsonLoads(text: string): JsonValue {
  let i = 0;
  const n = text.length;

  const ws = () => {
    while (i < n) {
      const c = text.charCodeAt(i);
      if (c === 0x20 || c === 0x09 || c === 0x0a || c === 0x0d) i += 1;
      else break;
    }
  };
  const fail = (what: string): never => {
    throw new JsonError(`${what} at ${i}`);
  };

  const string = (): string => {
    // text[i] is the opening quote
    i += 1;
    let out = '';
    for (;;) {
      if (i >= n) fail('unterminated string');
      const c = text.charCodeAt(i);
      if (c === 0x22) {
        i += 1;
        return out;
      }
      if (c === 0x5c) {
        const e = text[i + 1];
        i += 2;
        switch (e) {
          case '"': out += '"'; break;
          case '\\': out += '\\'; break;
          case '/': out += '/'; break;
          case 'b': out += '\b'; break;
          case 'f': out += '\f'; break;
          case 'n': out += '\n'; break;
          case 'r': out += '\r'; break;
          case 't': out += '\t'; break;
          case 'u': {
            const hex = text.slice(i, i + 4);
            if (!/^[0-9a-fA-F]{4}$/.test(hex)) fail('bad \\u escape');
            i += 4;
            out += String.fromCharCode(parseInt(hex, 16));
            break;
          }
          default:
            fail('bad escape');
        }
        continue;
      }
      if (c < 0x20) fail('control character in string');
      out += text[i];
      i += 1;
    }
  };

  const number = (): JsonNumber => {
    const m = /^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][-+]?\d+)?/.exec(text.slice(i));
    if (!m) fail('bad number');
    const raw = m![0];
    i += raw.length;
    return new JsonNumber(raw, Number(raw));
  };

  const value = (): JsonValue => {
    ws();
    if (i >= n) fail('unexpected end');
    const c = text[i]!;
    if (c === '{') {
      i += 1;
      const obj = new Map<string, JsonValue>();
      ws();
      if (text[i] === '}') {
        i += 1;
        return obj;
      }
      for (;;) {
        ws();
        if (text[i] !== '"') fail('expected property name');
        const key = string();
        ws();
        if (text[i] !== ':') fail('expected colon');
        i += 1;
        obj.set(key, value());
        ws();
        if (text[i] === ',') {
          i += 1;
          continue;
        }
        if (text[i] === '}') {
          i += 1;
          return obj;
        }
        fail('expected , or }');
      }
    }
    if (c === '[') {
      i += 1;
      const arr: JsonValue[] = [];
      ws();
      if (text[i] === ']') {
        i += 1;
        return arr;
      }
      for (;;) {
        arr.push(value());
        ws();
        if (text[i] === ',') {
          i += 1;
          continue;
        }
        if (text[i] === ']') {
          i += 1;
          return arr;
        }
        fail('expected , or ]');
      }
    }
    if (c === '"') return string();
    if (text.startsWith('true', i)) {
      i += 4;
      return true;
    }
    if (text.startsWith('false', i)) {
      i += 5;
      return false;
    }
    if (text.startsWith('null', i)) {
      i += 4;
      return null;
    }
    if (c === '-' || (c >= '0' && c <= '9')) return number();
    // NaN, Infinity, -Infinity and anything else: json.loads with parse_constant raising.
    return fail('unexpected token');
  };

  const result = value();
  ws();
  if (i < n) fail('extra data');
  return result;
}

/* ------------------------------------------------------------ Python repr */

/** Python's `str.isprintable()` for one code point. The ASCII space is the one printable separator. */
const NOT_PRINTABLE = /[\p{Cc}\p{Cf}\p{Cs}\p{Co}\p{Cn}\p{Zl}\p{Zp}\p{Zs}]/u;

function pyStrRepr(s: string): string {
  const quote = s.includes("'") && !s.includes('"') ? '"' : "'";
  let out = quote;
  for (const ch of s) {
    const cp = ch.codePointAt(0)!;
    if (ch === quote || ch === '\\') out += `\\${ch}`;
    else if (ch === '\n') out += '\\n';
    else if (ch === '\r') out += '\\r';
    else if (ch === '\t') out += '\\t';
    else if (cp !== 0x20 && NOT_PRINTABLE.test(ch)) {
      const hex = cp.toString(16);
      out += cp < 0x100 ? `\\x${hex.padStart(2, '0')}` : cp < 0x10000 ? `\\u${hex.padStart(4, '0')}` : `\\U${hex.padStart(8, '0')}`;
    } else out += ch;
  }
  return out + quote;
}

/** `repr(float(raw))`: shortest round-trip digits, exponent form outside 1e-4 ≤ |x| < 1e16. */
export function pyFloatRepr(x: number): string {
  if (Number.isNaN(x)) return 'nan';
  if (x === Infinity) return 'inf';
  if (x === -Infinity) return '-inf';
  if (x === 0) return Object.is(x, -0) ? '-0.0' : '0.0';
  const sign = x < 0 ? '-' : '';
  const [mantissa, expText] = Math.abs(x).toExponential().split('e');
  const digits = mantissa!.replace('.', '');
  const exp = Number(expText);
  if (exp < -4 || exp >= 16) {
    const m = digits.length > 1 ? `${digits[0]}.${digits.slice(1)}` : digits;
    const e = Math.abs(exp).toString().padStart(2, '0');
    return `${sign}${m}e${exp < 0 ? '-' : '+'}${e}`;
  }
  if (exp >= 0) {
    const whole = digits.length > exp + 1 ? digits.slice(0, exp + 1) : digits + '0'.repeat(exp + 1 - digits.length);
    const frac = digits.length > exp + 1 ? digits.slice(exp + 1) : '0';
    return `${sign}${whole}.${frac}`;
  }
  return `${sign}0.${'0'.repeat(-exp - 1)}${digits}`;
}

/** Python's `repr` of a value `json.loads` produced, or of a missing key. */
export function pyRepr(v: JsonValue | undefined): string {
  if (v === null || v === undefined) return 'None';
  if (v === true) return 'True';
  if (v === false) return 'False';
  if (typeof v === 'string') return pyStrRepr(v);
  if (v instanceof JsonNumber) return v.isInt ? BigInt(v.raw).toString() : pyFloatRepr(v.value);
  if (Array.isArray(v)) return `[${v.map(pyRepr).join(', ')}]`;
  return `{${[...v].map(([k, val]) => `${pyStrRepr(k)}: ${pyRepr(val)}`).join(', ')}}`;
}

/* --------------------------------------------------------------- extract */

/**
 * Python's `\s` for str patterns is `str.isspace()`, which is not JavaScript's
 * `\s`: Python includes U+001C–U+001F and U+0085, JavaScript includes U+FEFF.
 */
const PY_WS = '[\\t\\n\\x0b\\x0c\\r\\x1c-\\x1f \\x85\\xa0\\u1680\\u2000-\\u200a\\u2028\\u2029\\u202f\\u205f\\u3000]';
const FENCE = new RegExp(`\`\`\`json${PY_WS}*([\\s\\S]*?)\`\`\``);
const MONTH = /^\d{4}-\d{2}$/;

export function extractCandidate(text: string): string | null {
  const match = FENCE.exec(text);
  if (match) return match[1]!;
  const start = text.indexOf('{');
  if (start < 0) return null;
  let depth = 0;
  let inString = false;
  let escaped = false;
  for (let i = start; i < text.length; i += 1) {
    const ch = text[i];
    if (inString) {
      if (escaped) escaped = false;
      else if (ch === '\\') escaped = true;
      else if (ch === '"') inString = false;
      continue;
    }
    if (ch === '"') inString = true;
    else if (ch === '{') depth += 1;
    else if (ch === '}') {
      depth -= 1;
      if (depth === 0) return text.slice(start, i + 1);
    }
  }
  return null;
}

/** `[object, null]` or `[null, reason]`. Validation of fields is the engine's job. */
export function parseReply(text: string | null | undefined): [Map<string, JsonValue> | null, string | null] {
  if (!text) return [null, 'no JSON object found in the reply'];
  const candidate = extractCandidate(text);
  if (candidate === null) return [null, 'no JSON object found in the reply'];
  let value: JsonValue;
  try {
    value = pyJsonLoads(candidate);
  } catch {
    return [null, 'the JSON object could not be parsed'];
  }
  if (!isJsonObject(value)) return [null, 'the reply parsed but is not a JSON object'];
  return [value, null];
}

/**
 * Is `month` a YYYY-MM inside the window? Python's `\d` and `$` are looser
 * than JavaScript's, but the membership test that follows makes the regex
 * redundant either way: only a month that IS in the window passes.
 */
export function isWindowMonth(month: unknown, months: readonly string[]): month is string {
  return typeof month === 'string' && MONTH.test(month) && months.includes(month);
}

/* ------------------------------------------------------------------ emit */

/**
 * Python's `json.dumps(obj)` with its defaults: `", "` and `": "` separators,
 * `ensure_ascii=True`. The scripted policies emit their replies through this
 * so a recorded arm's reply strings are byte-identical on both sides.
 */
export function pyJsonDumps(value: unknown): string {
  if (value === null || value === undefined) return 'null';
  if (typeof value === 'string') {
    // ensure_ascii escapes everything above '~', DEL included, one \uXXXX per UTF-16 unit.
    return JSON.stringify(value).replace(/[\u007f-\uffff]/g, (c) => `\\u${c.charCodeAt(0).toString(16).padStart(4, '0')}`);
  }
  if (typeof value === 'boolean') return value ? 'true' : 'false';
  if (typeof value === 'number') {
    if (!Number.isInteger(value)) throw new Error(`pyJsonDumps: non-integer number ${value}`);
    return String(value);
  }
  if (Array.isArray(value)) return `[${value.map(pyJsonDumps).join(', ')}]`;
  const record = value as Record<string, unknown>;
  return `{${Object.keys(record)
    .filter((k) => record[k] !== undefined)
    .map((k) => `${pyJsonDumps(k)}: ${pyJsonDumps(record[k])}`)
    .join(', ')}}`;
}

/* ---------------------------------------------------------------- render */

const hours2 = (minutes: number) => (minutes / 60).toFixed(2);

export function systemPrompt(): string {
  const screen = hours2(SCREEN_MINUTES);
  const turn = hours2(TURN_MINUTES);
  const lookup = hours2(LOOKUP_MINUTES);
  const doc = hours2(DOC_MINUTES);
  return `You are a transaction-monitoring analyst at a bank. One alert has been assigned to you.

All data is synthetic. No real customer, account, transaction or counterparty exists.

Your screen shows the alert, the transactions that triggered it, the customer's KYC file and a twelve-month summary. Reading it has already cost ${screen} analyst-hours. You have ${MAX_TURNS} replies. Each reply must be exactly ONE JSON object (in a \`\`\`json fence, or as the first {...} in the reply), and every reply costs ${turn} hours whether or not it is accepted.

Lookups (each may be made once):
  {"action":"lookup","what":"history","month":"YYYY-MM"}       itemised transactions for a month in the window     +${lookup} h
  {"action":"lookup","what":"counterparty","id":"CP-n"}        name, type, jurisdiction, first seen, how many OTHER alerted accounts share it, its transactions     +${lookup} h
  {"action":"lookup","what":"prior_alerts"}                    prior alerts on this customer with disposition, analyst and rationale     +${lookup} h
  {"action":"lookup","what":"documents"}                       documents on file     +${doc} h

Dispositions (end the episode):
  {"action":"close","reason":R,"cites":[ids],"note":"..."}
      R is one of CONSISTENT_WITH_PROFILE, DOCUMENTED_SOURCE_OF_FUNDS, PREVIOUSLY_REVIEWED, RULE_ARTEFACT
  {"action":"escalate","typology":T,"cites":[ids],"narrative":"..."}
      T is one of STRUCTURING, FUNNEL_ACCOUNT, MONEY_MULE, RAPID_MOVEMENT, UNKNOWN

\`cites\` must be a non-empty list of ids you have actually been shown: kyc.<field>, T-<n>, CP-<n>, A-<n> (a prior alert), doc.<id>. Citing an id you have not been shown is rejected. The note or narrative is displayed and never scored.

A close on an alert that should have been escalated scores zero. An escalation on a clean customer scores half on the counterweight and zero on the rest. Cite the facts that decide the case, not everything you saw. If you reach your last reply without a disposition the alert stays in the queue and the episode scores zero.`;
}

const commas = (n: number) => {
  const s = String(n);
  let out = '';
  for (let i = 0; i < s.length; i += 1) {
    if (i > 0 && (s.length - i) % 3 === 0) out += ',';
    out += s[i];
  }
  return out;
};

export function money(cents: number, currency: string): string {
  const sign = cents < 0 ? '-' : '';
  const abs = Math.abs(cents);
  return `${sign}${commas(Math.floor(abs / 100))}.${String(abs % 100).padStart(2, '0')} ${currency}`;
}

const padEnd = (s: string, w: number) => s.padEnd(w);
const padStart = (s: string, w: number) => s.padStart(w);

function txnLine(t: Txn, currency: string): string {
  const where = t.branch ? `branch ${t.branch}` : `${t.cp}`;
  return `  ${padEnd(t.id, 6)} ${t.date}  ${padEnd(t.channel, 9)} ${padStart(money(t.amount, currency), 18)}  ${where}`;
}

export function renderScreen(screen: Screen): string {
  const cur = screen.currency;
  const c = screen.customer;
  const a = screen.alert;
  const alertBlock = [
    `ALERT ${a.id}  rule ${a.rule} — ${a.rule_text}`,
    `  fired ${a.fired}`,
    '  triggering transactions:',
    ...screen.triggering_transactions.map((t) => txnLine(t, cur)),
  ];
  const kycLines = [`KYC FILE  ${c.id}  ${c.name}`];
  const record = c as unknown as Record<string, string | number | boolean | undefined>;
  for (const key of Object.keys(record).sort()) {
    if (key === 'id' || key === 'name' || record[key] === undefined) continue;
    let value: string | number | boolean = record[key]!;
    if (key === 'stated_income' || key === 'expected_monthly_cash') value = money(value as number, cur);
    else if (typeof value === 'boolean') value = value ? 'True' : 'False';
    kycLines.push(`  kyc.${key} = ${value}`);
  }
  const summary = [
    `TWELVE-MONTH SUMMARY (${cur})`,
    `  ${padEnd('month', 8)}${padStart('cash in', 16)}${padStart('cash out', 16)}${padStart('wires in', 16)}${padStart('wires out', 16)}${padStart('ach in', 16)}${padStart('ach out', 16)}${padStart('txns', 6)}`,
  ];
  for (const r of screen.summary) {
    const cell = (v: number) => padStart(commas(Math.floor(v / 100)), 16);
    summary.push(
      `  ${padEnd(r.month, 8)}${cell(r.cash_in)}${cell(r.cash_out)}${cell(r.wires_in)}${cell(r.wires_out)}${cell(r.ach_in)}${cell(r.ach_out)}${padStart(String(r.count), 6)}`,
    );
  }
  const cps = ['COUNTERPARTIES on screen (top by twelve-month volume):'];
  for (const cp of screen.screen_counterparties) {
    cps.push(`  ${padEnd(cp.id, 6)} ${padEnd(cp.direction, 4)} ${padStart(String(cp.count), 3)} txns  ${money(cp.total, cur)}`);
  }
  const blocks = screen.screen_order === 'kyc_first' ? [kycLines, alertBlock] : [alertBlock, kycLines];
  blocks.push(summary, cps, [
    `Reporting threshold in this jurisdiction: ${money(screen.threshold, cur)}`,
    'SYNTHETIC DATA — no real customer, account or transaction.',
  ]);
  return blocks.map((b) => b.join('\n')).join('\n');
}

export type LookupResult =
  | { month: string; transactions: Txn[] }
  | Counterparty
  | { alerts: PriorAlert[] }
  | { documents: Document[] };

export function renderLookup(what: Lookup, result: LookupResult, currency: string): string {
  if (what === 'history') {
    const r = result as { month: string; transactions: Txn[] };
    return [`HISTORY ${r.month}: ${r.transactions.length} transactions`, ...r.transactions.map((t) => txnLine(t, currency))].join('\n');
  }
  if (what === 'counterparty') {
    const r = result as Counterparty;
    return [
      `COUNTERPARTY ${r.id}  ${r.name}`,
      `  type ${r.type}; jurisdiction ${r.jurisdiction}; first seen ${r.first_seen}`,
      `  other alerted accounts sharing this counterparty in twelve months: ${r.other_alerted_accounts}`,
      `  transactions: ${r.transactions.join(', ')}`,
    ].join('\n');
  }
  if (what === 'prior_alerts') {
    const r = result as { alerts: PriorAlert[] };
    if (!r.alerts.length) return 'PRIOR ALERTS: none on file.';
    const lines = ['PRIOR ALERTS:'];
    for (const p of r.alerts) {
      lines.push(`  ${p.id}  ${p.rule}  fired ${p.fired}  ${p.disposition} after ${p.elapsed_days} day(s) by ${p.analyst}`);
      lines.push(`         ${p.rationale}`);
    }
    return lines.join('\n');
  }
  const r = result as { documents: Document[] };
  return ['DOCUMENTS ON FILE:', ...r.documents.map((d) => `  ${d.id}  [${d.kind}]  ${d.description}`)].join('\n');
}

export function renderRejection(reason: string, turnsLeft: number): string {
  const plural = turnsLeft === 1 ? 'y' : 'ies';
  return `That reply was not accepted: ${reason}. It still cost a turn. You have ${turnsLeft} repl${plural} left.`;
}

export function renderCost(minutesSpent: number, turnsLeft: number): string {
  const plural = turnsLeft === 1 ? 'y' : 'ies';
  return `Hours spent so far: ${hours2(minutesSpent)}. You have ${turnsLeft} repl${plural} left.`;
}
