// Browser runtime for the compiled autocomplete model (zero dependencies, no DOM;
// works on the main thread, in a Web Worker, or in Bun/Node).
//
// web/model.bin holds the compiled transformer in sparse form (see export_web.py):
// the always-on unigram neuron plus, for every MLP neuron, its context tokens and
// the logit changes it writes. Evaluating only the neurons whose context matches
// gives the same next-token distribution as the dense model, without the zeros.
// The confidence policy mirrors suggest.py.
//
//   const ac = await CompiledAutocomplete.load("/model.bin");
//   ac.suggest("write a python fun").ghost   // "ction to"

import { decodePieces, preTokenize } from "../tokenizer.ts";

export type Policy = { tau: number; min_support: number; max_words: number };
export type Step = { text: string; prob: number; support: number };
export type Suggestion = { ghost: string; prob: number; steps: Step[] };

type IdArray = Uint16Array | Uint32Array;
type Order = {
  k: number;
  ctx: IdArray;
  support: Float32Array;
  offsets: Uint32Array;
  dids: IdArray;
  dvals: Float32Array;
  index: Map<string, number>;
};

const MAGIC = "TACB";

export class CompiledAutocomplete {
  readonly vocab: string[];
  readonly nPrev: number;
  readonly order: number;
  readonly eosId: number;
  readonly unkId: number;
  readonly policy: Policy;
  readonly bytes: number;
  readonly contexts: number;
  private text: string[];
  private pieceIndex: Map<string, number>;
  private base: Float32Array;
  private orders: Order[];
  private oov: Record<string, number>;
  private stopAfter: Set<string>;
  private candCache = new Map<string, Int32Array>();
  private distCache = new Map<string, { probs: Float64Array; support: number }>();

  /** Fetch (with optional progress callback) or wrap an already-loaded buffer. */
  static async load(source: string | ArrayBuffer, onProgress?: (fraction: number) => void): Promise<CompiledAutocomplete> {
    if (source instanceof ArrayBuffer) return new CompiledAutocomplete(source);
    const resp = await fetch(source);
    if (!resp.ok || !resp.body) throw new Error(`failed to fetch ${source}: ${resp.status}`);
    const total = Number(resp.headers.get("content-length")) || 0;
    const reader = resp.body.getReader();
    const chunks: Uint8Array[] = [];
    let received = 0;
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
      received += value.length;
      if (total) onProgress?.(received / total);
    }
    const bytes = new Uint8Array(received);
    let off = 0;
    for (const c of chunks) {
      bytes.set(c, off);
      off += c.length;
    }
    onProgress?.(1);
    return new CompiledAutocomplete(bytes.buffer);
  }

  constructor(buf: ArrayBuffer) {
    const view = new DataView(buf);
    const magic = new TextDecoder().decode(new Uint8Array(buf, 0, 4));
    if (magic !== MAGIC) throw new Error("not a tiny-autocomplete model.bin");
    const headerLen = view.getUint32(8, true);
    const header = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 12, headerLen)));
    const dataStart = 12 + headerLen + ((8 - ((12 + headerLen) % 8)) % 8);
    const sections = new Map<string, { offset: number; bytes: number }>();
    for (const s of header.sections) sections.set(s.name, s);
    const Ids = header.id_bytes === 2 ? Uint16Array : Uint32Array;
    const f32 = (name: string) => {
      const s = sections.get(name)!;
      return new Float32Array(buf, dataStart + s.offset, s.bytes / 4);
    };
    const u32 = (name: string) => {
      const s = sections.get(name)!;
      return new Uint32Array(buf, dataStart + s.offset, s.bytes / 4);
    };
    const ids = (name: string) => {
      const s = sections.get(name)!;
      return new Ids(buf, dataStart + s.offset, s.bytes / header.id_bytes);
    };

    this.vocab = header.vocab;
    this.text = this.vocab.map((p) => decodePieces(p));
    this.pieceIndex = new Map(this.vocab.map((p, i) => [p, i]));
    this.nPrev = header.n_prev;
    this.order = this.nPrev + 1;
    this.eosId = header.eos_id;
    this.unkId = header.unk_id;
    this.oov = header.oov_prefix_rates;
    this.policy = header.policy;
    this.stopAfter = new Set(header.stop_after);
    this.bytes = buf.byteLength;
    this.base = f32("base");
    this.orders = header.orders.map(({ k, contexts }: { k: number; contexts: number }) => {
      const ctx = ids(`ctx${k}`);
      const index = new Map<string, number>();
      for (let i = 0; i < contexts; i++) index.set(ctx.subarray(i * k, (i + 1) * k).join(","), i);
      return { k, ctx, support: f32(`support${k}`), offsets: u32(`offsets${k}`), dids: ids(`dids${k}`), dvals: f32(`dvals${k}`), index };
    });
    this.contexts = this.orders.reduce((n, o) => n + o.index.size, 0);
  }

  /** Last `order` token ids, oldest first; positions before 0 read BOS like the attention heads. */
  private window(ids: number[]): number[] {
    const t = ids.length - 1;
    const w: number[] = [];
    for (let a = this.nPrev; a >= 0; a--) w.push(ids[Math.max(t - a, 0)]);
    return w;
  }

  /** Next-token distribution and support of the longest matched context. */
  dist(ids: number[]): { probs: Float64Array; support: number } {
    const window = this.window(ids);
    const key = window.join(",");
    const hit = this.distCache.get(key);
    if (hit) return hit;
    const logits = Float64Array.from(this.base);
    let support = 0;
    for (const o of this.orders) {
      const i = o.index.get(window.slice(-o.k).join(","));
      if (i === undefined) continue;
      for (let j = o.offsets[i]; j < o.offsets[i + 1]; j++) logits[o.dids[j]] += o.dvals[j];
      support = o.support[i];
    }
    let max = -Infinity;
    for (const v of logits) if (v > max) max = v;
    let z = 0;
    const probs = new Float64Array(logits.length);
    for (let i = 0; i < logits.length; i++) z += probs[i] = Math.exp(logits[i] - max);
    for (let i = 0; i < probs.length; i++) probs[i] /= z;
    if (this.distCache.size > 2000) this.distCache.clear();
    const out = { probs, support };
    this.distCache.set(key, out);
    return out;
  }

  private candidates(partial: string): Int32Array {
    let c = this.candCache.get(partial);
    if (!c) {
      const out: number[] = [];
      for (let i = this.unkId + 1; i < this.vocab.length; i++) if (this.vocab[i].startsWith(partial)) out.push(i);
      c = Int32Array.from(out);
      if (this.candCache.size > 5000) this.candCache.clear();
      this.candCache.set(partial, c);
    }
    return c;
  }

  /** Greedy suggestion chain with cumulative probabilities, before thresholding. */
  steps(text: string, maxWords = this.policy.max_words): Step[] {
    const pcs = preTokenize(text);
    if (!pcs.length) return [];
    const partial = pcs[pcs.length - 1];
    const cands = this.candidates(partial);
    if (!cands.length) return [];
    const ids = [this.eosId, ...pcs.slice(0, -1).map((p) => this.pieceIndex.get(p) ?? this.unkId)];
    let { probs, support } = this.dist(ids);
    let best = cands[0];
    let sum = 0;
    for (const c of cands) {
      sum += probs[c];
      if (probs[c] > probs[best]) best = c;
    }
    let prob = (1 - (this.oov[partial] ?? 0)) * (probs[best] / sum);
    const out: Step[] = [{ text: this.text[best].slice(decodePieces(partial).length), prob, support }];
    ids.push(best);
    while (!this.stopAfter.has(this.vocab[ids[ids.length - 1]]) && out.length < maxWords) {
      ({ probs, support } = this.dist(ids));
      let nxt = 0;
      for (let i = 1; i < probs.length; i++) if (probs[i] > probs[nxt]) nxt = i;
      if (nxt === this.eosId) break;
      prob *= probs[nxt];
      out.push({ text: this.text[nxt], prob, support });
      ids.push(nxt);
    }
    return out;
  }

  static apply(steps: Step[], policy: Policy): { ghost: string; prob: number } {
    let ghost = "";
    let prob = 0;
    for (const st of steps.slice(0, policy.max_words)) {
      if (st.prob < policy.tau || st.support < policy.min_support) break;
      ghost += st.text;
      prob = st.prob;
    }
    return { ghost, prob };
  }

  /** Ghost text to show after `text`, or "" when the model isn't confident. */
  suggest(text: string, policy: Policy = this.policy): Suggestion {
    const steps = this.steps(text, policy.max_words);
    return { ...CompiledAutocomplete.apply(steps, policy), steps };
  }
}