// The Canopy reader on one screen, in Node: the exported ONNX model (scripts/export_reader.py) plus its
// tokenizer, encoded exactly as canopy/classifier/model.py:encode() and decoded as canopy/reader/evaluate.py.
// A reference for Daintree's integration; src/parity.ts checks it against the Python outputs.
import { existsSync, readFileSync } from "node:fs";
import { basename, dirname, join, resolve } from "node:path";
import * as ort from "onnxruntime-node";
import { AutoTokenizer, env } from "@huggingface/transformers";

export interface ReaderMeta {
  max_len: number;
  header_tokens: number;
  line_tokens: number;
  question_candidate_rows: number;
  status_candidate_rows: number;
  inputs: string[];
  outputs: string[];
  categories: string[];
  classes: Record<string, string[]>;
  thresholds: { summary_first: number; summary_more: number; option: number; gate: number; gate_categories: string[];
    second?: number; card?: string };
}

export interface Reading {
  category: string;
  category_probs: number[];
  needs_attention: number; // gated and rescaled: >= 0.5 means the card needs the user
  needs_attention_raw: number;
  blocked_on_user: number;
  notable: number;
  question_line: number; // line indices into the prepared screen, -1 for none
  status_line: number; // the headline
  action_line: number;
  task_line: number;
  summary_lines: number[]; // most important first
  card_lines: number[]; // the summary-only card: line one (the best line), then line two when it is worth showing
  card_mode: string; // "summary-only" for models trained for it
  option_lines: number[]; // screen order
  key_lines: number[];
  tests: string;
  changes: string;
  risk: string;
  instruction: string;
  progress: string;
  failure_repeats: string;
  score: string;
}

const cp = (s: string) => [...s].length; // Python len(): code points, not UTF-16 units
// Python's str.strip() whitespace (str.isspace), which differs from JS \s: it includes \x1c-\x1f and \x85 and
// excludes \ufeff. Training used Python, so blank-line skipping and candidate rows must use the same set.
const PY_WS = "\t\n\v\f\r\x1c\x1d\x1e\x1f \x85\xa0\u1680\u2000-\u200a\u2028\u2029\u202f\u205f\u3000";
const STRIP = new RegExp(`^[${PY_WS}]+|[${PY_WS}]+$`, "gu");
const strip = (s: string) => s.replace(STRIP, "");
const sigmoid = (x: number) => 1 / (1 + Math.exp(-x));

function softmax(xs: number[]): number[] {
  const m = Math.max(...xs);
  const e = xs.map((x) => Math.exp(x - m));
  const s = e.reduce((a, b) => a + b, 0);
  return e.map((x) => x / s);
}

export function lineCandidates(lines: string[], rows: number): number[] {
  const start = Math.max(0, lines.length - rows);
  const out: number[] = [];
  for (let i = start; i < lines.length; i++) if (cp(strip(lines[i])) > 3) out.push(i);
  return out;
}

export class CanopyReader {
  meta: ReaderMeta;
  session: ort.InferenceSession;
  tok: any;
  cls: number;
  sep: number;
  pad: number;

  private constructor(meta: ReaderMeta, session: ort.InferenceSession, tok: any) {
    this.meta = meta;
    this.session = session;
    this.tok = tok;
    this.cls = tok.model.tokens_to_ids.get(tok.cls_token ?? "[CLS]");
    this.sep = tok.model.tokens_to_ids.get(tok.sep_token ?? "[SEP]");
    this.pad = tok.model.tokens_to_ids.get(tok.pad_token ?? "[PAD]");
  }

  /** `dir` holds reader.json, the tokenizer files and model.<variant>.onnx. */
  static async load(dir: string, variant = "int8", threads = 0): Promise<CanopyReader> {
    // config.json first: Hugging Face counts a download per fetch of it (reader.json is the same file).
    const metaFile = existsSync(join(dir, "config.json")) ? join(dir, "config.json") : join(dir, "reader.json");
    const meta = JSON.parse(readFileSync(metaFile, "utf8")) as ReaderMeta;
    env.allowRemoteModels = false;
    env.localModelPath = dirname(resolve(dir)) + "/";
    const tok = await AutoTokenizer.from_pretrained(basename(resolve(dir)), { local_files_only: true } as any);
    const opts: ort.InferenceSession.SessionOptions = { graphOptimizationLevel: "all" };
    if (threads) opts.intraOpNumThreads = threads;
    const session = await ort.InferenceSession.create(join(dir, `model.${variant}.onnx`), opts);
    return new CanopyReader(meta, session, tok);
  }

  private ids(text: string): number[] {
    return this.tok.encode(text, { add_special_tokens: false });
  }

  /** As encode(): [CLS] state-json [SEP] then "[i] line\n" pieces filled from the bottom up, [SEP]. */
  /** `state` must hold integer numbers only (as Daintree sends): JSON.stringify writes 1.0 as 1, where the
   * Python training encoder wrote 1.0, so a float in the state would tokenize differently. */
  encode(state: object, lines: string[]) {
    const m = this.meta;
    const header = this.ids(JSON.stringify(state)).slice(0, m.header_tokens);
    let budget = m.max_len - header.length - 3;
    const pieces: [number, number[]][] = [];
    for (let i = lines.length - 1; i >= 0; i--) {
      if (!strip(lines[i])) continue;
      const ids = this.ids(`[${i}] ${lines[i]}\n`).slice(0, m.line_tokens);
      if (ids.length > budget) break;
      budget -= ids.length;
      pieces.push([i, ids]);
    }
    pieces.reverse();
    const ids = [this.cls, ...header, this.sep];
    const lineIndex: number[] = [];
    const linePos: number[] = [];
    for (const [i, p] of pieces) {
      lineIndex.push(i);
      linePos.push(ids.length);
      ids.push(...p);
    }
    ids.push(this.sep);
    const q = new Set(lineCandidates(lines, m.question_candidate_rows));
    const s = new Set(lineCandidates(lines, m.status_candidate_rows));
    return { ids, lineIndex, linePos, qMask: lineIndex.map((i) => q.has(i)), sMask: lineIndex.map((i) => s.has(i)) };
  }

  async read(state: object, lines: string[]): Promise<Reading> {
    const e = this.encode(state, lines);
    const n = e.ids.length;
    const k = Math.max(1, e.lineIndex.length);
    const pad = <T>(xs: T[], fill: T) => [...xs, ...Array(k - xs.length).fill(fill)];
    const feeds: Record<string, ort.Tensor> = {
      input_ids: new ort.Tensor("int64", BigInt64Array.from(e.ids.map(BigInt)), [1, n]),
      attention_mask: new ort.Tensor("int64", BigInt64Array.from(e.ids.map(() => 1n)), [1, n]),
      line_pos: new ort.Tensor("int64", BigInt64Array.from(pad(e.linePos, 0).map(BigInt)), [1, k]),
      line_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.lineIndex.map(() => 1), 0)), [1, k]),
      q_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.qMask.map(Number), 0)), [1, k]),
      s_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.sMask.map(Number), 0)), [1, k]),
    };
    const out = await this.session.run(feeds);
    const arr = (name: string) => Array.from(out[name].data as Float32Array, Number);
    return this.decode(arr, e.lineIndex, e.qMask);
  }

  decode(arr: (name: string) => number[], idx: number[], qMask: boolean[]): Reading {
    const m = idx.length;
    const t = this.meta.thresholds;
    const pick = (name: string) => {
      const scores = arr(name);
      const cand = [...scores.slice(0, m), scores[scores.length - 1]];
      let j = 0;
      for (let x = 1; x < cand.length; x++) if (cand[x] > cand[j]) j = x;
      return j === m ? -1 : idx[j];
    };
    const logit = arr("summary_lines").slice(0, m);
    const pSum = logit.map(sigmoid);
    const ranked = [...logit.keys()].sort((a, b) => logit[b] - logit[a] || a - b); // by logit: no saturation ties
    const summary: number[] = [];
    for (const j of ranked.slice(0, 3)) {
      if (pSum[j] >= (summary.length ? t.summary_more : t.summary_first)) summary.push(idx[j]);
      else break;
    }
    const first = pick("status_line");
    const cardLines = first >= 0 ? [first] : [];
    if (t.second != null) {
      const next = ranked.map((j) => idx[j]).filter((i, k) => i !== first && pSum[ranked[k]] >= (t.second as number));
      cardLines.push(...next.slice(0, 1));
    }
    const pOpt = arr("option_lines").slice(0, m).map(sigmoid);
    const key = arr("key_lines").slice(0, m);
    const probs = softmax(arr("category"));
    const raw = sigmoid(arr("needs_attention")[0]);
    const needs = t.gate_categories.reduce((a, c) => a + probs[this.meta.categories.indexOf(c)], 0);
    const g = Math.max(raw, needs);
    const argmax = (xs: number[]) => xs.reduce((b, x, i) => (x > xs[b] ? i : b), 0);
    const cls = (name: string) => this.meta.classes[name][argmax(arr(name))];
    return {
      category: this.meta.categories[argmax(probs)],
      category_probs: probs,
      needs_attention: g < t.gate ? (0.5 * g) / t.gate : 0.5 + (0.5 * (g - t.gate)) / (1 - t.gate),
      needs_attention_raw: raw,
      blocked_on_user: sigmoid(arr("blocked_on_user")[0]),
      notable: sigmoid(arr("notable")[0]),
      question_line: pick("question_line"),
      status_line: pick("status_line"),
      action_line: pick("action_line"),
      task_line: pick("task_line"),
      summary_lines: summary,
      card_lines: cardLines,
      card_mode: t.card ?? "headline",
      option_lines: idx.filter((_, j) => qMask[j] && pOpt[j] >= t.option),
      key_lines: idx.filter((_, j) => key[j] > 0),
      tests: cls("tests"),
      changes: cls("changes"),
      risk: cls("risk"),
      instruction: cls("instruction"),
      progress: cls("progress"),
      failure_repeats: cls("failure_repeats"),
      score: cls("score"),
    };
  }
}

// --- The card as Daintree shows it: line picks turned into clean text ---------------------------------

// Bullets, tree and box glyphs, spinners, prompt markers: stripped only when whitespace (or the end) follows,
// so `**bold**` and `>file` survive.
const LEAD = new RegExp(`^(?:[⏺•●✦⎿└├│┃║■✻✶✢✳✽✺·*❯›>▸▶◆◇-](?:[${PY_WS}]+|$))+`, "u");
const BOX_TAIL = new RegExp(`[${PY_WS}]*[│┃║][${PY_WS}]*$`, "u");
const TIMER = new RegExp(`[${PY_WS}]*\\((?:\\d[^)]*)?(?:esc to \\w+|thinking|thought for|ctrl\\+\\w to \\w+)[^)]*\\)[${PY_WS}]*$`, "iu");
const SENTENCE_END = /[.!?:;)"'`]$/u;

/** One screen line as card text: glyphs, borders and timers off, words kept. null when nothing is left. */
export function displayText(line: string | undefined): string | null {
  if (line == null) return null;
  const text = strip(strip(line).replace(LEAD, "").replace(BOX_TAIL, "").replace(TIMER, ""));
  return cp(text) >= 2 ? text : null;
}

export interface Card {
  category: string;
  needsAttention: boolean;
  score: number;
  headline: string | null;
  summary: string[]; // up to 3, most important first; a sentence wrapped over lines is one entry
  question: string | null;
  options: string[];
  action: string | null;
  task: string | null;
  tests: string;
  changes: string;
  progress: number | null;
}

/** Does screen line `j` carry on the sentence of the line above it? Yes when it was picked with it, or
 * when it has no bullet of its own and starts in lowercase or indented (an agent's own line break). */
/** A one-character fragment displayText() drops for length (a lone "1"), not text the cleanup removed. */
function short(line: string): string | null {
  const t = strip(line);
  return cp(t) === 1 ? t : null;
}

/** Does line `cur` carry on the sentence of `prev` (the agent wrapped it onto a new row)? Mirrors
 * canopy_reader.py `_continues`. */
function continuesFrom(prev: string, cur: string): boolean {
  if (!strip(prev) || !strip(cur) || LEAD.test(strip(cur))) return false;
  const pt = displayText(prev);
  if (!pt || SENTENCE_END.test(pt)) return false;
  // Python: re.match(r"^\s{2,}\S"), str.islower() on the first character (Unicode Lowercase), a bracket or quote.
  return new RegExp(`^[${PY_WS}]{2,}[^${PY_WS}]`, "u").test(cur) || /^\p{Lowercase}/u.test(strip(cur)) || /^[(\[{'"`]/u.test(strip(cur));
}

/** Line `i` as a whole sentence: back to where a wrapped sentence starts and on to where it ends, up to four
 * rows each way. A new bullet is always a new item. */
function sentenceAt(lines: string[], i: number, picks: number[], used: Set<number>): string | null {
  let k = i;
  while (k > 0 && i - k < 4 && !used.has(k - 1) && continuesFrom(lines[k - 1], lines[k])) k--;
  for (let x = k; x <= i; x++) used.add(x);
  let text: string | null = lines.slice(k, i + 1).map((x, n) => displayText(x) ?? (k + n !== i ? short(x) : null))
    .filter((x): x is string => !!x).join(" ") || null;
  for (let j = i + 1; text && j <= i + 4 && j < lines.length && !SENTENCE_END.test(text) && !used.has(j); j++) {
    const nxt = lines[j];
    if (!strip(nxt) || LEAD.test(strip(nxt))) break;
    if (!picks.includes(j) && !continuesFrom(lines[j - 1], nxt)) break;
    used.add(j);
    const more = displayText(nxt) ?? short(nxt); // keep a one-character fragment ("1") the sentence carries on into
    if (more) text += " " + more;
  }
  return text;
}

/** Summary picks as text, most important first, each a whole sentence where it wraps. */
function summaryText(lines: string[], picks: number[], used: Set<number>): string[] {
  const out: string[] = [];
  for (const i of picks) {
    if (used.has(i)) continue;
    const text = sentenceAt(lines, i, picks, used);
    if (text) out.push(text);
  }
  return out;
}

export function card(r: Reading, lines: string[]): Card {
  const asks = r.category === "approval" || r.category === "question";
  const line = (i: number) => (i >= 0 ? displayText(lines[i]) : null);
  const used = new Set<number>();
  if (r.card_mode === "summary-only") {
    // No headline (the terminal title is the title): up to two lines, line one the best on screen.
    return {
      category: r.category, needsAttention: r.needs_attention >= 0.5, score: Number(r.score), headline: null,
      summary: summaryText(lines, r.card_lines.length === 2 && Math.abs(r.card_lines[0] - r.card_lines[1]) <= 3
        ? [...r.card_lines].sort((a, b) => a - b) : r.card_lines, used),
      question: asks ? line(r.question_line) : null,
      options: asks ? r.option_lines.map((i) => displayText(lines[i])).filter((x): x is string => !!x) : [],
      action: r.category === "approval" ? line(r.action_line) : null,
      task: line(r.task_line), tests: r.tests, changes: r.changes,
      progress: r.progress === "none" ? null : Number(r.progress),
    };
  }
  const headline = r.status_line >= 0 ? sentenceAt(lines, r.status_line, r.summary_lines, used) : null;
  return {
    category: r.category,
    needsAttention: r.needs_attention >= 0.5,
    score: Number(r.score),
    headline,
    summary: summaryText(lines, r.summary_lines, used),
    question: asks ? line(r.question_line) : null,
    options: asks ? r.option_lines.map((i) => displayText(lines[i])).filter((x): x is string => !!x) : [],
    action: r.category === "approval" ? line(r.action_line) : null,
    task: line(r.task_line),
    tests: r.tests,
    changes: r.changes,
    progress: r.progress === "none" ? null : Number(r.progress),
  };
}

/** The card's summary text for given picks, as card() builds it: for parity tests against canopy_reader.py. */
export function summaryTextForTesting(lines: string[], picks: number[]): string[] {
  const p = picks.length === 2 && Math.abs(picks[0] - picks[1]) <= 3 ? [...picks].sort((a, b) => a - b) : picks;
  return summaryText(lines, p, new Set());
}