import type { Model, Progress, Tokenizer, Transformers } from "./runtime-types.ts";
import type { Generation, WorkerRequest, WorkerResponse } from "./types.ts";
// Tokenling brain: a module worker that runs the fine-tuned GPT-2 (small, 4-bit)
// with transformers.js on WebAssembly. Weights are fetched once from this site
// and kept in the browser's Cache Storage; nothing is sent anywhere.

let T: Transformers,
  tokenizer: Tokenizer,
  model: Model,
  busy = false;
const CTX = 1024;

const post = <K extends WorkerResponse["type"]>(
  type: K,
  data: Omit<Extract<WorkerResponse, { type: K }>, "type">,
) => self.postMessage({ type, ...data });

async function load({
  runtime,
  models,
  modelId,
  dtype,
  threads,
}: Extract<WorkerRequest, { type: "load" }>) {
  T = await import(`${runtime}/transformers.min.js`);
  if (threads) T.env.backends.onnx.wasm.numThreads = threads;
  T.env.allowRemoteModels = false;
  T.env.allowLocalModels = true;
  T.env.localModelPath = models;
  T.env.useBrowserCache = true;
  T.env.backends.onnx.wasm.wasmPaths = `${runtime}/`;
  T.env.backends.onnx.logLevel = "error";
  const files = new Map<string, [number, number]>();
  const progress_callback = (p: Progress) => {
    if (p.status === "progress" && p.total) {
      files.set(p.file, [p.loaded, p.total]);
      let loaded = 0,
        total = 0;
      for (const [l, t] of files.values()) {
        loaded += l;
        total += t;
      }
      post("progress", { loaded, total, file: p.file });
    } else if (p.status === "done") {
      post("file", { file: p.file });
    }
  };
  tokenizer = await T.AutoTokenizer.from_pretrained(modelId, { progress_callback });
  // The model stores 4-bit weights behind DequantizeLinear nodes.
  // disable_quant_qdq makes ONNX Runtime fold them to fp32 at load instead of
  // fusing them into MatMulNBits, which is about 5x slower in WebAssembly
  // (train/quantize.py); without it the model does not load in a usable way.
  // disable_prepacking skips a second packed copy of every weight (-160 MB).
  model = await T.AutoModelForCausalLM.from_pretrained(modelId, {
    device: "wasm",
    dtype,
    progress_callback,
    session_options: { extra: { session: { disable_quant_qdq: "1", disable_prepacking: "1" } } },
  });
  post("loading", { stage: "warming up" });
  await model.generate({ ...tokenizer("hi"), max_new_tokens: 1, do_sample: false });
}

// Stop at end of line (chat) or right after a lookup tag so the app can
// append its {result}; diaries run until <|endoftext|>.
// GPT-2 has a single "][" token, so a lookup can close mid-token.
const LOOKUP_END = /\[(clock|look|yesterday|recall:[^\]\n]{1,80})\]\[?$/;

function makeStopper(
  promptLen: number,
  mode: Generation["mode"],
  state: { text: string; reason: string },
) {
  class Stop extends T.StoppingCriteria {
    _call(input_ids: (number | bigint)[][]) {
      const ids = input_ids[0].slice(promptLen).map(Number);
      const text = tokenizer.decode(ids, { skip_special_tokens: true });
      state.text = text;
      if (mode === "line") {
        if (text.includes("\n")) {
          state.reason = "line";
          return [true];
        }
        if (LOOKUP_END.test(text)) {
          state.reason = "lookup";
          return [true];
        }
      } else if (/\n##|\n\n\n/.test(text)) {
        state.reason = "end";
        return [true];
      }
      return [false];
    }
  }
  return new Stop();
}

// lines: newest last. Drops the oldest lines until header+lines+prefix fit in
// ctx tokens (prompt plus new tokens). Idle turns pass a smaller ctx because
// prefill dominates their cost.
function fit(header: string, lines: string[], prefix: string, maxNew: number, ctx = CTX) {
  const budget = Math.min(ctx, CTX) - maxNew - 4;
  const count = (s: string) => tokenizer.encode(s).length;
  const hTok = count(header) + count(prefix);
  const kept = [];
  let used = hTok;
  for (let i = lines.length - 1; i >= 0; i--) {
    const n = count(lines[i] + "\n");
    if (used + n > budget) break;
    kept.unshift(lines[i]);
    used += n;
  }
  return header + kept.map((l) => l + "\n").join("") + prefix;
}

async function generate({
  id,
  header,
  lines,
  prefix,
  mode,
  maxNew,
  temperature,
  ctx,
}: Extract<WorkerRequest, { type: "generate" }>) {
  const prompt = fit(header, lines, prefix, maxNew, ctx);
  const inputs = tokenizer(prompt);
  const promptLen = inputs.input_ids.dims[1];
  const state = { text: "", reason: "length" };
  let first = 0;
  const streamer = new T.TextStreamer(tokenizer, {
    skip_prompt: true,
    callback_function: () => {
      first ||= performance.now();
      post("token", { id, text: state.text });
    },
  });
  const t0 = performance.now();
  await model.generate({
    ...inputs,
    max_new_tokens: maxNew,
    do_sample: true,
    temperature,
    top_k: 40,
    top_p: 0.92,
    stopping_criteria: [makeStopper(promptLen, mode, state)],
    streamer,
  });
  let text = state.text;
  if (state.reason === "line") text = text.slice(0, text.indexOf("\n"));
  if (state.reason === "lookup") text = text.replace(/\]\[$/, "]");
  if (mode !== "line")
    text = text.replace(/\n##[\s\S]*$/, "").replace(/<\|endoftext\|>[\s\S]*$/, "");
  // prefillMs: prompt processing up to the first new token; tps: decoding speed after it.
  const end = performance.now(),
    n = tokenizer.encode(text).length;
  const prefillMs = (first || end) - t0;
  post("done", {
    id,
    text,
    prompt,
    reason: state.reason,
    promptTokens: promptLen,
    prefillMs,
    tps: n > 1 && first ? (n - 1) / ((end - first) / 1000) : 0,
  });
}

self.onmessage = async ({ data }: MessageEvent<WorkerRequest>) => {
  try {
    if (data.type === "load") {
      await load(data);
      // The runtime files this worker fetched, so the page can keep them for offline use.
      post("ready", {
        dtype: data.dtype,
        resources: performance.getEntriesByType("resource").map((e) => e.name),
      });
    } else if (data.type === "generate") {
      if (busy) return post("error", { id: data.id, message: "busy" });
      busy = true;
      try {
        await generate(data);
      } finally {
        busy = false;
      }
    }
  } catch (err) {
    post("error", {
      id: data.type === "generate" ? data.id : undefined,
      message: err instanceof Error ? err.message : String(err),
    });
  }
};