// The engine lives here, not on the page's thread.
//
// candle in WASM is synchronous and CPU-bound: constructing the model, encoding
// a sentence, and sampling a token all occupy their thread outright. Run that
// where the UI lives and the tab stops responding, including the progress text
// that was supposed to explain the wait. Everything heavy happens in this
// worker and the page only ever handles messages.

import init, { Engine, Loader } from "./blackwindow_wasm.js";

let engine = null;

const post = (type, payload) => self.postMessage({ type, ...payload });

const DB = "blackwindow", STORE = "blobs", CHUNK = 32 * 1024 * 1024;

// IndexedDB rather than the Cache API. Cache.put rejects the 380MB model with
// an opaque internal error, and storing it as one 380MB value fails the same
// way, so it goes in as 32MB chunks under a manifest.
function idb() {
  return new Promise((res, rej) => {
    const r = indexedDB.open(DB, 1);
    r.onupgradeneeded = () => r.result.createObjectStore(STORE);
    r.onsuccess = () => res(r.result);
    r.onerror = () => rej(r.error);
  });
}

function tx(db, mode, fn) {
  return new Promise((res, rej) => {
    const t = db.transaction(STORE, mode);
    const out = fn(t.objectStore(STORE));
    // "result" in out, not out.result !== undefined: a miss is a request whose
    // result is undefined, and the looser test returned the request itself,
    // which is truthy and made every miss look like a corrupt hit.
    t.oncomplete = () => res(out && typeof out === "object" && "result" in out ? out.result : out);
    t.onerror = () => rej(t.error);
  });
}

async function cacheGet(db, key, label, sink) {
  const meta = await tx(db, "readonly", (s) => s.get(`${key}:meta`));
  if (!meta) return null;
  const out = sink ? null : new Uint8Array(meta.size);
  let off = 0;
  for (let i = 0; i < meta.chunks; i++) {
    const part = await tx(db, "readonly", (s) => s.get(`${key}:${i}`));
    if (!part) return null;
    const u = new Uint8Array(part);
    if (sink) sink(u);
    else out.set(u, off);
    off += u.length;
    // Reading 382MB back from IndexedDB is as expensive as downloading it, and
    // reporting nothing here once made a stalled counter look like a ceiling.
    post("progress", { label, got: off, total: meta.size });
  }
  if (off !== meta.size) return null;
  return sink ? true : out;
}

async function cachePut(db, key, bytes) {
  const chunks = Math.ceil(bytes.length / CHUNK);
  for (let i = 0; i < chunks; i++) {
    const slice = bytes.slice(i * CHUNK, (i + 1) * CHUNK);
    await tx(db, "readwrite", (s) => s.put(slice.buffer, `${key}:${i}`));
  }
  await tx(db, "readwrite", (s) =>
    s.put({ size: bytes.length, chunks }, `${key}:meta`));
}

// The cache key is the URL, so bumping a versioned filename strands the old
// blob forever. Drop anything the current build no longer asks for.
async function pruneCache(keep) {
  try {
    const db = await idb();
    const keys = await tx(db, "readonly", (s) => s.getAllKeys());
    const live = new Set(keep);
    const stale = keys.filter((k) => !live.has(
      String(k).replace(/:(meta|\d+)$/, "").replace(/:blob$/, "")));
    if (!stale.length) return;
    await tx(db, "readwrite", (s) => { for (const k of stale) s.delete(k); });
    post("log", { line: `cache: dropped ${stale.length} entries from an older build` });
  } catch (e) {
    post("log", { line: `cache prune skipped (${e})` });
  }
}

async function fetchBytes(url, label) {
  // Cache anything served from an absolute URL. Those are the big immutable
  // blobs on a CDN, and the URL is the cache key, so a versioned filename
  // invalidates correctly on rebuild. Relative paths are small local files.
  const cacheable = /^https?:/.test(url);
  let db = null;
  if (cacheable) {
    try {
      db = await idb();
      const hit = await cacheGet(db, url);
      if (hit) {
        post("log", { line: `${label}: from cache (${(hit.length / 1e6).toFixed(0)} MB)` });
        return hit;
      }
    } catch (e) {
      post("log", { line: `${label}: cache unavailable (${e})` });
    }
  }

  const r = await fetch(url);
  if (!r.ok) throw new Error(`${label}: HTTP ${r.status} for ${url}`);
  const total = +r.headers.get("content-length") || 0;
  const reader = r.body.getReader();

  // Preallocate when the length is known and fill in place. Accumulating
  // chunks and concatenating afterwards holds the file twice at once, which
  // is what puts a 382MB model over an iOS tab's memory limit.
  let out = total ? new Uint8Array(total) : null;
  const chunks = out ? null : [];
  let got = 0, lastPost = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    if (out) {
      if (got + value.length > total) {
        throw new Error(`${label}: server sent more than ${total} bytes`);
      }
      out.set(value, got);
    } else {
      chunks.push(value);
    }
    got += value.length;
    if (got - lastPost > 4e6) {
      lastPost = got;
      post("progress", { label, got, total });
    }
  }
  post("progress", { label, got, total: total || got });
  if (!out) {
    out = new Uint8Array(got);
    let o = 0;
    for (const c of chunks) { out.set(c, o); o += c.length; }
  } else if (got !== total) {
    throw new Error(`${label}: expected ${total} bytes, got ${got}`);
  }

  if (db) {
    try {
      await cachePut(db, url, out);
      post("log", { line: `${label}: cached for next time` });
    } catch (e) {
      post("log", { line: `${label}: not cached (${e})` });
    }
  }
  return out;
}

// Stream one blob straight into the loader, caching as it goes. Nothing here
// ever holds more than one CHUNK, which is the entire point: the previous
// version assembled the whole 382MB file in JS and then handed it to wasm,
// which meant three copies existed at once and a phone refused the third.
async function stageBlob(url, label, loader) {
  let db = null;
  if (/^https?:/.test(url)) {
    try { db = await idb(); } catch (e) { post("log", { line: `${label}: cache unavailable (${e})` }); }
  }

  if (db) {
    const meta = await tx(db, "readonly", (s) => s.get(`${url}:meta`));
    if (meta) {
      loader.begin(meta.size);
      const ok = await cacheGet(db, url, label, (u) => loader.push(u));
      if (ok) {
        post("log", { line: `${label}: from cache (${(meta.size / 1e6).toFixed(0)} MB)` });
        return;
      }
      // A tab killed mid-write leaves chunks without a usable manifest. Drop
      // them, or they sit there forever and the refetch pays for them twice.
      post("log", { line: `${label}: cache entry incomplete, refetching` });
      try {
        await tx(db, "readwrite", (s) => {
          for (let i = 0; i < (meta.chunks || 0); i++) s.delete(`${url}:${i}`);
          s.delete(`${url}:meta`);
        });
      } catch (e) { /* best effort */ }
      loader.begin(0);
    }
  }

  const r = await fetch(url);
  if (!r.ok) throw new Error(`${label}: HTTP ${r.status} for ${url}`);
  const total = +r.headers.get("content-length") || 0;
  loader.begin(total);

  const reader = r.body.getReader();
  const buf = new Uint8Array(CHUNK);
  let used = 0, chunkIdx = 0, got = 0, lastPost = 0;

  const flush = async () => {
    if (!used) return;
    loader.push(buf.subarray(0, used));
    if (db) {
      try {
        const copy = buf.slice(0, used);
        await tx(db, "readwrite", (s) => s.put(copy.buffer, `${url}:${chunkIdx}`));
      } catch (e) {
        db = null;
        post("log", { line: `${label}: not cached (${e})` });
      }
    }
    chunkIdx++;
    used = 0;
  };

  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    let off = 0;
    while (off < value.length) {
      const n = Math.min(CHUNK - used, value.length - off);
      buf.set(value.subarray(off, off + n), used);
      used += n; off += n; got += n;
      if (used === CHUNK) await flush();
    }
    if (got - lastPost > 4e6) {
      lastPost = got;
      post("progress", { label, got, total });
    }
  }
  await flush();
  post("progress", { label, got, total: total || got });

  if (db) {
    try {
      await tx(db, "readwrite", (s) => s.put({ size: got, chunks: chunkIdx }, `${url}:meta`));
      post("log", { line: `${label}: cached for next time` });
    } catch (e) {
      post("log", { line: `${label}: not cached (${e})` });
    }
  }
}

// The model is kept as a Blob, never as bytes.
//
// Everything else is small enough to stage inside wasm, but the GGUF is not:
// candle copies it into tensors, so a staged copy means the file exists twice,
// and wasm memory never shrinks. Blobs are backed by the browser's disk store,
// and FileReaderSync gives a worker synchronous random access to one, which is
// exactly what a Rust Read + Seek needs.
async function modelBlob(url, label) {
  let db = null;
  try { db = await idb(); } catch (e) { post("log", { line: `${label}: cache unavailable (${e})` }); }
  const key = `${url}:blob`;

  if (db) {
    const hit = await tx(db, "readonly", (s) => s.get(key));
    if (hit && hit.size) {
      post("log", { line: `${label}: from cache (${(hit.size / 1e6).toFixed(0)} MB)` });
      post("progress", { label, got: hit.size, total: hit.size });
      return hit;
    }
  }

  // response.blob() lets the browser stream straight into its own store. Doing
  // this by hand means the chunks pass through the JS heap first, which is the
  // copy we are here to avoid. The cost is a progress bar that only moves at
  // the end.
  post("log", { line: `${label}: downloading` });
  post("progress", { label, got: 0, total: 0 });
  const r = await fetch(url);
  if (!r.ok) throw new Error(`${label}: HTTP ${r.status} for ${url}`);
  const blob = await r.blob();
  post("progress", { label, got: blob.size, total: blob.size });

  if (db) {
    try {
      await tx(db, "readwrite", (s) => s.put(blob, key));
      post("log", { line: `${label}: cached for next time` });
    } catch (e) {
      post("log", { line: `${label}: not cached (${e})` });
    }
  }
  return blob;
}

async function load({ gguf, tokenizer, head, index, anchor, tapLayer, meanPool, stream }) {
  await init();
  post("log", { line: "fetching model, tokenizer, head, gallery" });
  await pruneCache([gguf, tokenizer, head, index, anchor].filter((u) => /^https?:/.test(u)));

  // Model first: it is the largest, so its peak must not overlap the gallery's.
  const loader = new Loader();
  await stageBlob(tokenizer, "tokenizer", loader);
  loader.finish_tokenizer();

  // Streaming is the default: it keeps the GGUF out of linear memory, which is
  // what lets an iPhone build the model at all. Staging is the fallback,
  // because a browser that cannot hand back a 382MB Blob should still work.
  let built = false;
  if (stream !== false) {
    try {
      const blob = await modelBlob(gguf, "model");
      post("log", { line: `building model (streaming ${(blob.size / 1e6).toFixed(0)} MB)` });
      const fr = new FileReaderSync();
      loader.finish_model_streaming(
        (off, len) => new Uint8Array(fr.readAsArrayBuffer(blob.slice(off, off + len))),
        blob.size,
      );
      built = true;
    } catch (e) {
      post("log", { line: `streaming load failed (${e.message || e}), staging instead` });
    }
  }
  if (!built) {
    await stageBlob(gguf, "model", loader);
    post("log", { line: "building model" });
    loader.finish_model();
  }

  await stageBlob(head, "head", loader);
  loader.finish_head();
  await stageBlob(index, "gallery", loader);
  loader.finish_index();

  post("log", { line: "building engine" });
  engine = loader.build(tapLayer, !!meanPool);

  const a = await fetchBytes(anchor, "anchor");
  const bytes = engine.apply_anchor(a);
  post("log", { line: `runtime anchor applied: ${(bytes / 1024).toFixed(1)} KB` });

  post("ready", { gallery: engine.gallery_size() });
}

self.onmessage = async (ev) => {
  const m = ev.data;
  try {
    switch (m.type) {
      case "load":
        await load(m);
        break;

      case "verifyAnchor": {
        // Confirms the shipped constant actually describes this browser.
        const cos = engine.anchor_drift(m.probes);
        post("anchorDrift", { cos });
        break;
      }

      case "chat": {
        const answer = engine.chat(m.text, m.maxTokens, m.temperature, m.k, m.seed,
          (piece, hits) => {
            if (hits) post("hits", { hits });
            if (piece) post("token", { piece });
            return true;
          });
        post("answer", { answer });
        break;
      }

      case "search":
        post("hits", { hits: engine.search(m.text, m.k), tag: m.tag });
        break;

      case "buildAxis": {
        engine.build_axis(m.name, m.positive, m.negative);
        const control = engine.add_random_control(m.name, 7);
        post("axisReady", { name: m.name, control });
        break;
      }

      case "steer": {
        const hits = engine.search_steered(m.text, m.axis, m.alpha, m.k);
        const retention = engine.retention(m.axis, m.alpha, [m.text]);
        post("steered", { hits, retention, label: m.label });
        break;
      }

      case "recalibrateHere": {
        const bytes = engine.recalibrate(m.anchors);
        post("log", { line: `re-derived anchor here: ${(bytes / 1024).toFixed(1)} KB` });
        post("ready", { gallery: engine.gallery_size() });
        break;
      }
    }
  } catch (e) {
    post("error", { message: String(e && e.message ? e.message : e) });
  }
};