// SPDX-License-Identifier: Apache-2.0
// Embedding worker: loads a decision-tuned EmbeddingGemma 2 (ONNX, same graph layout as
// onnx-community/embeddinggemma-2-ONNX) with transformers.js and returns L2-normalized
// 768-d embeddings (mean pooling over the attention mask, then L2 normalization).
// The model files are served next to this page (./models/...), so nothing is fetched from elsewhere.
import {
  AutoConfig,
  AutoModel,
  AutoProcessor,
  RawImage,
  RawVideo,
  RawVideoFrame,
  env,
} from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.3.1/dist/transformers.min.js";

const MODEL_ID = "egemma2-decide";
env.allowLocalModels = true;
env.allowRemoteModels = false;
env.localModelPath = new URL("./models/", self.location.href).pathname; // same-origin path, e.g. "/models/"
if (env.backends?.onnx?.wasm) env.backends.onnx.wasm.numThreads = 1; // static hosting is not cross-origin isolated

let processor = null;
let model = null;
let loaded = null; // { device, dtype, audio }
let decision = null; // { scale } from decision.json
let queue = Promise.resolve();

const post = (msg, transfer) => self.postMessage(msg, transfer || []);
const state = (phase, label, extra = {}) => post({ type: "state", phase, label, ...extra });

let gpuF16 = false;
async function hasWebGPU() {
  try {
    if (!("gpu" in navigator)) return false;
    const adapter = await navigator.gpu.requestAdapter();
    gpuF16 = !!adapter?.features?.has("shader-f16");
    return !!adapter;
  } catch {
    return false;
  }
}

async function loadModel({ device = "auto", dtype = "q4", audio = false }) {
  if (loaded && (device === "auto" || device === loaded.device) && dtype === loaded.requested && (!audio || loaded.audio)) {
    return loaded;
  }
  if (model) {
    try { await model.dispose(); } catch {}
    model = null;
    loaded = null;
  }
  let candidates;
  if (device === "auto") candidates = (await hasWebGPU()) ? ["webgpu", "wasm"] : ["wasm"];
  else candidates = [device];

  if (!decision) {
    try { decision = await (await fetch(`${env.localModelPath}${MODEL_ID}/decision.json`)).json(); } catch { decision = {}; }
  }
  if (!processor) {
    state("loading", "Loading processor");
    processor = await AutoProcessor.from_pretrained(MODEL_ID);
  }
  let lastErr = null;
  const requested = dtype;
  if (candidates.includes("webgpu") && device !== "auto") await hasWebGPU();
  for (const dev of candidates) {
    // Files shipped with this page (Space storage is limited to 1 GB):
    //   q4: 4-bit weights with plain 8-bit gathers for the embedding and position tables, so the same files
    //       run on WebGPU and on the WASM (CPU) runtime, which lacks GatherBlockQuantized;
    //   q4f16: 4-bit weights with fp16 math (WebGPU with shader-f16 only);
    //   audio: one 4-bit audio encoder used with both.
    const d = dev === "wasm" ? "q4" : requested === "auto" ? (gpuF16 ? "q4f16" : "q4") : requested;
    dtype = audio ? { model: d, vision_encoder: d, audio_encoder: "q4" } : d;
    try {
      const config = await AutoConfig.from_pretrained(MODEL_ID);
      if (!audio) config.audio_config = null; // text + vision only (video uses the vision encoder)
      state("loading", `Downloading model (${dev}${audio ? ", with audio encoder" : ""})`, { progress: 0, device: dev });
      const t0 = performance.now();
      model = await AutoModel.from_pretrained(MODEL_ID, {
        config,
        device: dev,
        dtype,
        progress_callback: (p) => {
          if (p.status === "progress_total") state("loading", `Downloading model (${dev})`, { progress: p.progress, device: dev, loadedBytes: p.loaded, totalBytes: p.total });
        },
      });
      const sessionMs = performance.now() - t0;
      state("loading", "Warming up", { progress: 100, device: dev });
      const tw = performance.now();
      const warm = await processor(["task: classification | query: warm up"]);
      const out = await model(warm);
      out.sentence_embedding.dispose?.();
      const warmMs = performance.now() - tw;
      const label = typeof dtype === "string" ? dtype : Object.entries(dtype).map(([k, v]) => `${k.replace("_encoder", "")} ${v}`).join(", ");
      loaded = { device: dev, dtype: label, requested, audio, loadMs: performance.now() - t0, sessionMs, warmMs, scale: decision?.scale || null };
      state("ready", `Ready on ${dev.toUpperCase()} (${label})`, { device: dev, audio });
      return loaded;
    } catch (e) {
      lastErr = e;
      console.warn(`Loading on ${dev} failed`, e);
      if (model) { try { await model.dispose(); } catch {} model = null; }
      state("loading", `Loading on ${dev} failed, trying next backend`, { device: dev });
    }
  }
  state("error", `Could not load the model: ${lastErr?.message || lastErr}`);
  throw lastErr;
}

function rows(tensor) {
  const [n, d] = tensor.dims;
  const data = tensor.data;
  const out = [];
  for (let i = 0; i < n; i++) {
    const v = new Float32Array(d);
    let s = 0;
    for (let j = 0; j < d; j++) { v[j] = data[i * d + j]; s += v[j] * v[j]; }
    s = Math.sqrt(s) || 1;
    for (let j = 0; j < d; j++) v[j] /= s;
    out.push(v);
  }
  return out;
}

let lastModelMs = 0;
async function run(inputs) {
  const tm = performance.now();
  const out = await model(inputs);
  lastModelMs = performance.now() - tm;
  const emb = rows(out.sentence_embedding);
  try { out.sentence_embedding.dispose?.(); } catch {}
  return emb;
}

// texts: array of strings, embedded in small batches (right padding; padded tokens are masked out)
async function embedTexts(texts) {
  const res = [];
  const B = 8;
  for (let i = 0; i < texts.length; i += B) {
    const batch = texts.slice(i, i + B);
    res.push(...(await run(await processor(batch))));
  }
  return res;
}

// One input that may interleave text with an image, audio clip or video.
// If the text already holds a placeholder (<|image|>, <|audio|>, <|video|>) it is used as is;
// otherwise the placeholder is put first, as in the trained prompt format. Empty text embeds the media alone.
async function embedInput(input) {
  const { kind } = input;
  const text = (input.text || "").trim();
  const withPh = (ph) => (text.includes(ph) ? text : `${ph}${text}`);
  if (kind === "text" || !kind) {
    return (await run(await processor([text])))[0];
  }
  if (kind === "image" || kind === "rgba") {
    if (processor.image_processor) processor.image_processor.max_soft_tokens = input.budget || 280;
    const img = kind === "rgba"
      ? new RawImage(new Uint8ClampedArray(input.data), input.width, input.height, 4)
      : await RawImage.fromBlob(input.blob);
    if (!text) return (await run(await processor(null, img)))[0];
    return (await run(await processor([withPh("<|image|>")], [[img]])))[0];
  }
  if (kind === "audio") {
    if (!loaded?.audio) throw new Error("The audio encoder is not loaded.");
    const audio = input.audio instanceof Float32Array ? input.audio : new Float32Array(input.audio);
    if (!text) return (await run(await processor(null, null, audio)))[0];
    return (await run(await processor([withPh("<|audio|>")], null, [audio])))[0];
  }
  if (kind === "video") {
    if (processor.image_processor) processor.image_processor.max_soft_tokens = input.budget || 280;
    const frames = input.frames.map((f) => new RawVideoFrame(new RawImage(new Uint8ClampedArray(f.data), f.width, f.height, 4), f.timestamp));
    const video = new RawVideo(frames, input.duration);
    if (!text) return (await run(await processor(null, null, null, video)))[0];
    return (await run(await processor([withPh("<|video|>")], null, null, [video])))[0];
  }
  throw new Error(`Unknown input kind: ${kind}`);
}

async function handle(msg) {
  const { id, op } = msg;
  try {
    if (op === "load") {
      const info = await loadModel(msg.options || {});
      post({ type: "result", id, result: info });
      return;
    }
    if (!model) throw new Error("Model is not loaded yet.");
    const t0 = performance.now();
    if (op === "embedTexts") {
      const embs = await embedTexts(msg.texts);
      post({ type: "result", id, result: { embeddings: embs, ms: performance.now() - t0 } }, embs.map((e) => e.buffer));
    } else if (op === "embedInput") {
      const emb = await embedInput(msg.input);
      post({ type: "result", id, result: { embedding: emb, ms: performance.now() - t0, modelMs: lastModelMs } }, [emb.buffer]);
    } else if (op === "embedInputs") {
      // several prompts over the same media (one pass each); used by live cards in the trained format
      const embs = [];
      for (const inp of msg.inputs) embs.push(await embedInput(inp));
      post({ type: "result", id, result: { embeddings: embs, ms: performance.now() - t0 } }, embs.map((e) => e.buffer));
    } else {
      throw new Error(`Unknown op ${op}`);
    }
  } catch (e) {
    post({ type: "error", id, message: e?.message || String(e) });
  }
}

self.addEventListener("message", ({ data }) => {
  // serialize all requests: one model, one queue
  queue = queue.then(() => handle(data));
});