/**
 * One episode, in the browser: the free screen, up to eight replies, and the
 * hour meter.
 *
 * This is the entry point of the port of `envs/alert_triage/alert_triage/` —
 * the generator, the protocol, the reward and the scripted analysts live in
 * the sibling modules and are re-exported here. "Port" is meant strictly:
 * `scripts/conformance.mjs` computes the three digests in
 * `envs/alert_triage/CONFORMANCE.txt` from this module and fails the build
 * if any differs from what the Python computed. That gate is what lets the
 * page claim it verified a recorded run rather than merely replayed it.
 *
 * Two discipline rules, enforced here rather than in the grader:
 *
 *   - **Shown ids.** A disposition may cite only ids the engine has actually
 *     rendered — KYC fields and triggering transactions from the screen, plus
 *     whatever lookups returned.
 *   - **Abort versus truncate.** `aborted` is exactly MAX_TURNS turns spent
 *     with no accepted disposition, and scores 0.000. Anything that stops
 *     earlier with no disposition is `truncated` and is not scored at all.
 *
 * Keep this module pure and dependency-free. It runs on the main thread and
 * under `node --test`; hashing is the conformance script's business.
 */

import { MAX_TURNS, canonicalJson, generate, screenOf } from './generator';
import type { Counterparty, Document, Label, PriorAlert, Screen, Txn, World } from './generator';
import { ACTIONS, CLOSE_REASONS, LOOKUPS, TYPOLOGIES, isWindowMonth, parseReply, pyRepr } from './protocol';
import type { JsonValue, Lookup, LookupResult } from './protocol';
import { POLICIES, SHIPPED } from './policies';
import type { LookupRecord, Policy, View } from './policies';
import { DOC_MINUTES, LOOKUP_MINUTES, SCREEN_MINUTES, TURN_MINUTES, scoreExact } from './score';
import type { Episode } from './score';

export * from './rng';
export * from './generator';
export * from './protocol';
export * from './score';
export * from './policies';

const LOOKUP_PRICE: Record<Lookup, number> = {
  history: LOOKUP_MINUTES,
  counterparty: LOOKUP_MINUTES,
  prior_alerts: LOOKUP_MINUTES,
  documents: DOC_MINUTES,
};

/** The canonical form of an accepted action. */
export type EngineAction =
  | { action: 'lookup'; what: 'history'; month: string }
  | { action: 'lookup'; what: 'counterparty'; id: string }
  | { action: 'lookup'; what: 'prior_alerts' | 'documents' }
  | { action: 'close'; reason: string; cites: string[]; note: string }
  | { action: 'escalate'; typology: string; cites: string[]; narrative: string };

export type Disposition = Extract<EngineAction, { action: 'close' | 'escalate' }>;

export interface Step {
  accepted: boolean;
  action: EngineAction | null;
  /** Why a rejected reply was rejected. */
  reason: string | null;
  /** What an accepted lookup returned. */
  result: LookupResult | null;
  minutes_spent: number;
  turns_spent: number;
  done: boolean;
}

export type Outcome = 'solved' | 'failed' | 'aborted';

export class Engine {
  readonly world: World;
  turns = 0;
  rejected = 0;
  minutes = SCREEN_MINUTES;
  actions: EngineAction[] = [];
  lookups: LookupRecord[] = [];
  disposition: Disposition | null = null;
  shown = new Set<string>();
  lastRejection: string | null = null;
  private readonly screenCache: Screen;

  constructor(world: World) {
    this.world = world;
    this.screenCache = this.buildScreen();
    const screen = this.screenCache;
    for (const k of Object.keys(screen.customer)) {
      if (k !== 'id' && k !== 'name' && (screen.customer as unknown as Record<string, unknown>)[k] !== undefined) this.shown.add(`kyc.${k}`);
    }
    this.shown.add(screen.alert.id);
    for (const t of screen.triggering_transactions) {
      this.shown.add(t.id);
      if (t.cp) this.shown.add(t.cp);
    }
    for (const cp of screen.screen_counterparties) this.shown.add(cp.id);
  }

  /* ------------------------------------------------------------- views */

  private buildScreen(): Screen {
    const screen = screenOf(this.world);
    if (screen.screen_order === 'kyc_first') screen.summary = [...screen.summary].reverse();
    return screen;
  }

  screen(): Screen {
    return this.screenCache;
  }

  /**
   * Everything the model has been shown, and nothing it has not. The scripted
   * policies take this and only this: no label, typology, planted set,
   * overlay or reference is in it.
   */
  view(): View {
    return {
      screen: this.screenCache,
      lookups: [...this.lookups],
      turns_spent: this.turns,
      turns_left: MAX_TURNS - this.turns,
      rejected: this.rejected,
      minutes_spent: this.minutes,
      last_rejection: this.lastRejection,
    };
  }

  get done(): boolean {
    return this.disposition !== null || this.turns >= MAX_TURNS;
  }

  get aborted(): boolean {
    return this.disposition === null && this.turns >= MAX_TURNS;
  }

  /* -------------------------------------------------------------- step */

  step(reply: string | null | undefined): Step {
    if (this.done) {
      return { accepted: false, action: null, reason: 'the episode is over', result: null, minutes_spent: this.minutes, turns_spent: this.turns, done: true };
    }
    const [obj, parseReason] = parseReply(reply);
    let action: EngineAction | null = null;
    let reason = parseReason;
    if (obj !== null) [action, reason] = this.validate(obj);
    this.turns += 1;
    this.minutes += TURN_MINUTES;
    if (action === null) {
      this.rejected += 1;
      this.lastRejection = reason;
      return { accepted: false, action: null, reason, result: null, minutes_spent: this.minutes, turns_spent: this.turns, done: this.done };
    }

    this.lastRejection = null;
    this.actions.push(action);
    let result: LookupResult | null = null;
    if (action.action === 'lookup') {
      this.minutes += LOOKUP_PRICE[action.what];
      result = this.lookup(action);
      const key = action.what === 'history' ? action.month : action.what === 'counterparty' ? action.id : null;
      this.lookups.push({ what: action.what, key, result });
    } else {
      this.disposition = action;
    }
    return { accepted: true, action, reason: null, result, minutes_spent: this.minutes, turns_spent: this.turns, done: this.done };
  }

  private validate(obj: Map<string, JsonValue>): [EngineAction | null, string | null] {
    // A missing key is Python's `.get()` None; `pyRepr` prints both as None.
    const get = (k: string): JsonValue | undefined => obj.get(k);
    const action = get('action');
    if (!(ACTIONS as readonly unknown[]).includes(action)) {
      return [null, `unknown action ${pyRepr(action)}; use lookup, close or escalate`];
    }
    if (action === 'lookup') {
      const what = get('what');
      if (!(LOOKUPS as readonly unknown[]).includes(what)) {
        return [null, `unknown lookup ${pyRepr(what)}; use history, counterparty, prior_alerts or documents`];
      }
      const lookup = what as Lookup;
      if (lookup === 'history') {
        const month = get('month');
        if (!isWindowMonth(month, this.world.months)) {
          return [null, `month ${pyRepr(month)} is not in the twelve-month window`];
        }
        if (this.lookups.some((l) => l.what === 'history' && l.key === month)) {
          return [null, `history for ${month} has already been looked up`];
        }
        return [{ action: 'lookup', what: 'history', month }, null];
      }
      if (lookup === 'counterparty') {
        const cid = get('id');
        if (typeof cid !== 'string' || !this.shown.has(cid) || !cid.startsWith('CP-')) {
          return [null, `counterparty ${pyRepr(cid)} has not appeared on your screen or in a fetched month`];
        }
        if (this.lookups.some((l) => l.what === 'counterparty' && l.key === cid)) {
          return [null, `counterparty ${cid} has already been looked up`];
        }
        return [{ action: 'lookup', what: 'counterparty', id: cid }, null];
      }
      if (this.lookups.some((l) => l.what === lookup)) {
        return [null, `${lookup} has already been looked up`];
      }
      return [{ action: 'lookup', what: lookup }, null];
    }

    const cites = get('cites');
    if (!Array.isArray(cites) || cites.length === 0) return [null, 'cites must be a non-empty list of ids'];
    const seen: string[] = [];
    for (const c of cites) {
      if (typeof c !== 'string') return [null, 'cites must be a non-empty list of ids'];
      if (!seen.includes(c)) seen.push(c);
    }
    for (const c of seen) {
      if (!this.shown.has(c)) return [null, `cites an id you have not been shown: ${c}`];
    }
    if (action === 'close') {
      const reason = get('reason');
      if (!(CLOSE_REASONS as readonly unknown[]).includes(reason)) return [null, `unknown close reason ${pyRepr(reason)}`];
      const note = get('note');
      return [{ action: 'close', reason: reason as string, cites: seen, note: typeof note === 'string' ? note : '' }, null];
    }
    const typology = get('typology');
    if (!(TYPOLOGIES as readonly unknown[]).includes(typology)) return [null, `unknown typology ${pyRepr(typology)}`];
    const narrative = get('narrative');
    return [{ action: 'escalate', typology: typology as string, cites: seen, narrative: typeof narrative === 'string' ? narrative : '' }, null];
  }

  private lookup(action: Extract<EngineAction, { action: 'lookup' }>): LookupResult {
    const w = this.world;
    if (action.what === 'history') {
      const rows: Txn[] = w.transactions.filter((t) => t.month === action.month);
      for (const t of rows) {
        this.shown.add(t.id);
        if (t.cp) this.shown.add(t.cp);
      }
      return { month: action.month, transactions: rows };
    }
    if (action.what === 'counterparty') {
      const cp = w.counterparties.find((c) => c.id === action.id)!;
      for (const id of cp.transactions) this.shown.add(id);
      const copy: Counterparty = { ...cp };
      return copy;
    }
    if (action.what === 'prior_alerts') {
      for (const p of w.prior_alerts) this.shown.add(p.id);
      const alerts: PriorAlert[] = [...w.prior_alerts];
      return { alerts };
    }
    for (const d of w.documents) this.shown.add(d.id);
    const documents: Document[] = [...w.documents];
    return { documents };
  }

  /* ----------------------------------------------------------- outcome */

  /** solved | failed | aborted, or null while running / if truncated. */
  get outcome(): Outcome | null {
    if (this.disposition !== null) {
      const label: Label = this.world.label;
      const right = (label === 'benign') === (this.disposition.action === 'close');
      return right ? 'solved' : 'failed';
    }
    return this.aborted ? 'aborted' : null;
  }

  episode(truncated: boolean | null = null): Episode {
    const d = this.disposition;
    return {
      label: this.world.label,
      planted: this.world.planted,
      disposition: d === null ? null : d.action,
      cites: d === null ? [] : [...d.cites],
      minutes_spent: this.minutes,
      reference_minutes: this.world.reference_minutes ?? null,
      turns_spent: this.turns,
      rejected: this.rejected,
      truncated: truncated === null ? !this.done : truncated,
      typology: d === null ? null : d.action === 'escalate' ? d.typology : null,
      true_typology: this.world.typology ?? null,
      lookups: this.lookups.length,
    };
  }
}

/* ------------------------------------------------------------ reference */

/** Play a world to the end under one policy. Returns the finished Engine. */
export function runPolicy(policy: Policy, world: World): Engine {
  const engine = new Engine(world);
  while (!engine.done) engine.step(policy(engine.view()));
  return engine;
}

/**
 * The cheapest shipped policy that reaches the correct disposition, in
 * minutes. `[null, null]` if none does — which the Python's tests assert
 * never happens for seeds 0–4095.
 */
export function referenceFor(world: World): [number | null, string | null] {
  let best: [number, string] | null = null;
  for (const name of SHIPPED) {
    const engine = runPolicy(POLICIES[name]!, world);
    if (engine.outcome === 'solved' && (best === null || engine.minutes < best[0])) best = [engine.minutes, name];
  }
  return best ?? [null, null];
}

/**
 * The world with `reference_minutes` attached — part of the digest.
 *
 * The policies read only the engine's view, never the world, so attaching
 * their result here is not circular; it is the same move as wordle's
 * `reference_depth`, made part of the world so the browser recomputes it
 * instead of trusting the recorded number.
 */
export function worldForSeed(seed: number): World {
  const world = generate(seed);
  const [minutes, policy] = referenceFor(world);
  world.reference_minutes = minutes;
  world.reference_policy = policy;
  return world;
}

/** Re-score a recorded run from its reply strings alone. What `verify` does. */
export function replay(seed: number, replies: readonly (string | null)[], world: World | null = null): { engine: Engine; episode: Episode } {
  const engine = new Engine(world ?? worldForSeed(seed));
  for (const reply of replies) {
    if (engine.done) break;
    engine.step(reply);
  }
  return { engine, episode: engine.episode() };
}

/* ---------------------------------------------------------- conformance */
// The three digests, as the lines that get hashed. `conformance.py` is the
// definition; these produce the same byte strings for SHA-256 to consume. No
// hashing here so the module stays browser-clean.

const GRID_LABELS = ['benign', 'suspicious'] as const;
const GRID_DISPOSITIONS = ['close', 'escalate', 'none', 'truncated'] as const;
const GRID_CITES: string[][] = [['a'], ['a', 'b'], ['a', 'b', 'c', 'd'], ['x'], ['a', 'x'], ['a', 'b', 'x', 'y', 'z'], ['a', 'a', 'b'], ['kyc.z']];
const GRID_PLANTED: string[][][] = [[['a']], [['a', 'b']], [['a', 'b'], ['a']], [['c'], ['a', 'b', 'c', 'd']]];
const GRID_MINUTES = [33, 42, 60, 96, 150, 204];
const GRID_REFERENCE: (number | null)[] = [33, 42, 96, null];

/** Digest (a): one canonical world per seed, `reference_minutes` included. */
export function* worldLines(seeds: Iterable<number>): Generator<string> {
  for (const seed of seeds) yield canonicalJson(worldForSeed(seed));
}

/** Digest (b): `scoreExact` over the fixed grid, in `conformance.py`'s loop order. */
export function* scorerLines(): Generator<string> {
  for (const label of GRID_LABELS) {
    for (const disposition of GRID_DISPOSITIONS) {
      for (const cites of GRID_CITES) {
        for (const planted of GRID_PLANTED) {
          for (const minutes of GRID_MINUTES) {
            for (const reference of GRID_REFERENCE) {
              const none = disposition === 'none' || disposition === 'truncated';
              const ep: Episode = {
                label,
                planted,
                disposition: none ? null : disposition,
                cites: none ? [] : cites,
                minutes_spent: minutes,
                reference_minutes: reference,
                turns_spent: 1,
                rejected: 0,
                truncated: disposition === 'truncated',
              };
              yield canonicalJson(scoreExact(ep));
            }
          }
        }
      }
    }
  }
}

export interface CorpusCase {
  seed: number;
  note?: string;
  replies: (string | null)[];
}

/** Digest (c): parse + engine over the reply corpus on fixed worlds. */
export function protocolTrace(corpus: readonly CorpusCase[]): Record<string, unknown>[] {
  const out: Record<string, unknown>[] = [];
  for (const c of corpus) {
    const engine = new Engine(worldForSeed(c.seed));
    const turns: Record<string, unknown>[] = [];
    for (const reply of c.replies) {
      if (engine.done) break;
      const step = engine.step(reply);
      turns.push({ accepted: step.accepted, action: step.action, reason: step.reason, minutes: step.minutes_spent, turns: step.turns_spent });
    }
    const ep = engine.episode();
    out.push({ seed: c.seed, turns, outcome: engine.outcome, truncated: ep.truncated, score: scoreExact(ep) });
  }
  return out;
}

export function* protocolLines(corpus: readonly CorpusCase[]): Generator<string> {
  for (const entry of protocolTrace(corpus)) yield canonicalJson(entry);
}
