importScripts("vendor/tf.min.js", "vendor/tf-backend-wasm.js", "model.js");
let model, running = false, stopped = false;

async function load(backend, selection = "small") {
  const directory = selection === "large" ? "model/large/" : "model/";
  if (backend === "auto") {
    const adapter = await self.navigator.gpu?.requestAdapter().catch(() => null);
    const software = adapter?.isFallbackAdapter || adapter?.info?.isFallbackAdapter ||
      /swiftshader|software/i.test(adapter?.info?.description || "");
    backend = adapter && !software ? "webgpu" : "wasm";
  }
  if (backend === "webgpu") {
    try {
      importScripts("vendor/tf-backend-webgpu.js");
      if (!await tf.setBackend("webgpu")) throw new Error("WebGPU unavailable");
      await tf.ready();
    } catch (error) {
      backend = "wasm";
    }
  }
  if (backend === "wasm") {
    tf.wasm.setWasmPaths("vendor/");
    tf.wasm.setThreadsCount(1);
    await tf.setBackend("wasm");
    await tf.ready();
  }
  const config = await (await fetch(directory + "model-int8.json", {cache: "no-cache"})).json();
  const response = await fetch(directory + config.weight_file);
  if (!response.ok) throw new Error("Model download failed");
  const reader = response.body.getReader();
  const bytes = new Uint8Array(config.weight_bytes);
  let loaded = 0;
  for (;;) {
    const {value, done} = await reader.read();
    if (done) break;
    bytes.set(value, loaded);
    loaded += value.length;
    postMessage({type: "loading", loaded, total: config.weight_bytes});
  }
  if (loaded !== bytes.length) throw new Error("Incomplete model download");
  config.directory = directory;
  model = new TinyLM(config, tf.io.decodeWeights(bytes.buffer, config.weights));
  const warmup = model.step([32]);
  await warmup.logits.data();
  tf.dispose(warmup);
  postMessage({type: "ready", backend: tf.getBackend(), config});
}

function random(seed) {
  return () => {
    seed |= 0; seed = seed + 0x6D2B79F5 | 0;
    let t = Math.imul(seed ^ seed >>> 15, 1 | seed);
    t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t;
    return ((t ^ t >>> 14) >>> 0) / 4294967296;
  };
}

function choose(logits, temperature, rng) {
  const maximum = Math.max(...logits);
  const probabilities = Array.from(logits, x => Math.exp((x - maximum) / temperature));
  let target = rng() * probabilities.reduce((a, b) => a + b, 0);
  for (let i = 0; i < probabilities.length; i++) {
    target -= probabilities[i];
    if (target <= 0) return i;
  }
  return 255;
}

async function generate({prompt, count, temperature, seed}) {
  const input = Array.from(new TextEncoder().encode(prompt));
  if (!input.length || input.length >= model.config.context) throw new Error("Prompt must use 1–511 UTF-8 bytes");
  const limit = Math.min(count, model.config.context - input.length);
  stopped = false;
  let state;
  const start = performance.now();
  try {
    state = model.step(input);
    let logits = await state.logits.data();
    const prefillMs = performance.now() - start;
    const decoding = performance.now();
    const decoder = new TextDecoder();
    const rng = random(seed);
    let text = prompt;
    for (let i = 0; i < limit && !stopped; i++) {
      const byte = choose(logits, temperature, rng);
      text += decoder.decode(new Uint8Array([byte]), {stream: true});
      const positions = state.cache[0].key.shape[1];
      postMessage({type: "token", text, generated: i + 1, positions, prefillMs,
        bytesPerSecond: i ? i * 1000 / (performance.now() - decoding) : 0,
        cacheBytes: positions * model.config.layers * 2 * model.config.width * 4});
      if (i + 1 < limit && !stopped) {
        const previous = state;
        state = model.step([byte], previous.cache);
        tf.dispose(previous);
        logits = await state.logits.data();
      }
      if (i % 4 === 0) await new Promise(resolve => setTimeout(resolve, 0));
    }
    text += decoder.decode();
    postMessage({type: "done", text, stopped});
  } finally {
    if (state) tf.dispose(state);
  }
}

async function verify() {
  const {cases} = await (await fetch(model.config.directory + model.config.verification_file, {cache: "no-cache"})).json();
  let maximumError = 0, checks = 0;
  const baseline = tf.memory().numTensors;
  for (const test of cases) {
    let state;
    try {
      for (const reference of test.steps) {
        const previous = state;
        state = model.step(reference.input, previous ? previous.cache : []);
        if (previous) tf.dispose(previous);
        const actual = await state.logits.data();
        for (let i = 0; i < actual.length; i++) {
          if (!Number.isFinite(actual[i])) throw new Error("Non-finite browser logits");
          maximumError = Math.max(maximumError, Math.abs(actual[i] - reference.logits[i]));
        }
        checks++;
      }
    } finally {
      if (state) tf.dispose(state);
    }
  }
  const leaked = tf.memory().numTensors - baseline;
  if (maximumError > 0.002 || leaked) throw new Error(`Verification failed: error=${maximumError}, tensors=${leaked}`);
  postMessage({type: "verified", checks, maximumError, leaked, backend: tf.getBackend()});
}

self.onmessage = async ({data}) => {
  if (data.type === "stop") { stopped = true; return; }
  if (running) return;
  running = true;
  try {
    if (data.type === "load") await load(data.backend || "auto", data.model);
    if (data.type === "generate") await generate(data);
    if (data.type === "verify") await verify();
  } catch (error) {
    postMessage({type: "error", message: error.message});
  } finally {
    running = false;
  }
};