import * as ort from "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/ort.webgpu.bundle.min.mjs";
import { predict } from "./predict.js";

const REPO = "https://huggingface.co/mrfakename/Kumo-Tabular-ONNX/resolve/main/onnx";
ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/";
ort.env.wasm.numThreads = self.crossOriginIsolated ? Math.min(4, navigator.hardwareConcurrency || 1) : 1;

const sessions = new Map();
const post = (type, extra) => self.postMessage({ type, ...extra });

async function fetchModel(url, label) {
  let cache = null;
  try { cache = await caches.open("kumo-tabular-onnx-v1"); } catch {}
  const hit = cache && (await cache.match(url));
  if (hit) return new Uint8Array(await hit.arrayBuffer());
  const res = await fetch(url);
  if (!res.ok) throw new Error(`Download failed (${res.status}) for ${label}`);
  const total = +res.headers.get("content-length") || 0;
  const reader = res.body.getReader();
  const parts = []; let got = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    parts.push(value); got += value.length;
    post("status", { message: `Downloading ${label}… ${(got / 1e6).toFixed(0)}${total ? ` / ${(total / 1e6).toFixed(0)}` : ""} MB`, progress: total ? got / total : null });
  }
  const buf = new Uint8Array(got); let o = 0;
  for (const p of parts) { buf.set(p, o); o += p.length; }
  if (cache) { try { await cache.put(url, new Response(buf)); } catch {} }
  return buf;
}

async function getSession(size, task, backend) {
  const key = `${size}_${task}_${backend}`;
  if (sessions.has(key)) return sessions.get(key);
  // keep at most one session alive to bound memory
  for (const [k, s] of sessions) { try { await s.release(); } catch {} sessions.delete(k); }
  const label = `${size} ${task} model`;
  const bytes = await fetchModel(`${REPO}/${size}_${task}.onnx`, label);
  post("status", { message: `Initializing ${label} on ${backend.toUpperCase()}…`, progress: null });
  const providers = backend === "webgpu" ? ["webgpu", "wasm"] : ["wasm"];
  const session = await ort.InferenceSession.create(bytes, { executionProviders: providers, graphOptimizationLevel: "all" });
  sessions.set(key, session);
  return session;
}

self.onmessage = async ({ data }) => {
  if (data.type !== "predict") return;
  let backend = data.backend;
  if (backend === "webgpu" && !(navigator.gpu && (await navigator.gpu.requestAdapter().catch(() => null)))) backend = "wasm";
  try {
    const t0 = performance.now();
    const res = await predict({
      ort, ctx: data.ctx, qry: data.qry, target: data.target, task: data.task,
      numEstimators: data.numEstimators,
      getSession: (task) => getSession(data.size, task, backend),
      onProgress: (message, progress) => post("status", { message, progress }),
    });
    const secs = ((performance.now() - t0) / 1000).toFixed(1);
    post("result", {
      ...res,
      summary: `${res.task} · ${data.size} model on ${backend.toUpperCase()} · ${res.nContext} context rows · ${res.nQuery} predicted · ${res.nFeatures} features · ${secs}s`,
    });
  } catch (e) {
    console.error(e);
    post("error", { message: `Error: ${e.message || e}` });
  }
};