// Runs the MotionBricks G1 ONNX model with onnxruntime-web (WASM) off the main thread.
/* global ort */
importScripts('https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/ort.wasm.min.js');

let session = null;

async function cachedFetch(url, onProgress) {
  let cache = null;
  try { cache = await caches.open('motionbricks-models-v1'); } catch (e) { cache = null; }
  if (cache) {
    const hit = await cache.match(url);
    if (hit) {
      onProgress(1, 1, true);
      return new Uint8Array(await hit.arrayBuffer());
    }
  }
  const response = await fetch(url, { mode: 'cors' });
  if (!response.ok) throw new Error(`HTTP ${response.status} for ${url}`);
  const total = Number(response.headers.get('content-length')) || 0;
  const reader = response.body.getReader();
  const chunks = [];
  let received = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    received += value.length;
    onProgress(received, total, false);
  }
  const bytes = new Uint8Array(received);
  let offset = 0;
  for (const c of chunks) { bytes.set(c, offset); offset += c.length; }
  if (cache) {
    try { await cache.put(url, new Response(bytes, { headers: { 'content-type': 'application/octet-stream' } })); } catch (e) { /* quota */ }
  }
  return bytes;
}

self.onmessage = async (event) => {
  const msg = event.data;
  if (msg.type === 'load') {
    try {
      const threads = self.crossOriginIsolated ? Math.max(1, Math.min(8, (navigator.hardwareConcurrency || 4) - 1)) : 1;
      ort.env.wasm.numThreads = threads;
      ort.env.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/';
      let bytes = null;
      let lastError = null;
      for (const url of msg.urls) {
        try {
          bytes = await cachedFetch(url, (r, t, cached) => self.postMessage({ type: 'progress', received: r, total: t, cached }));
          break;
        } catch (e) { lastError = e; }
      }
      if (!bytes) throw lastError || new Error('Model download failed');
      self.postMessage({ type: 'compiling' });
      const t0 = performance.now();
      session = await ort.InferenceSession.create(bytes, { executionProviders: ['wasm'], graphOptimizationLevel: 'all' });
      self.postMessage({ type: 'ready', threads, compileMs: performance.now() - t0 });
    } catch (e) {
      self.postMessage({ type: 'error', message: String(e && e.message ? e.message : e) });
    }
    return;
  }
  if (msg.type === 'run') {
    const f = msg.feeds;
    const n = f.n || 1;
    const t0 = performance.now();
    try {
      const out = await session.run({
        context_qpos: new ort.Tensor('float32', f.context_qpos, [n, 4, 36]),
        target_qpos: new ort.Tensor('float32', f.target_qpos, [n, 4, 36]),
        has_global_root_values: new ort.Tensor('int32', f.has_global_root_values, [n, 8]),
        has_local_root_values: new ort.Tensor('int32', f.has_local_root_values, [n, 8]),
        has_local_poses: new ort.Tensor('int32', f.has_local_poses, [n, 8]),
        allowed_pred_num_tokens: new ort.Tensor('int32', f.allowed_pred_num_tokens, [n, 11]),
      });
      const pred_qpos = out.pred_qpos.data, pred_foot_contacts = out.pred_foot_contacts.data;
      const ntok = Array.from(out.pred_num_tokens.data, Number);
      self.postMessage({ type: 'result', id: msg.id, pred_qpos, pred_foot_contacts, ntok, ms: performance.now() - t0 },
        [pred_qpos.buffer, pred_foot_contacts.buffer]);
    } catch (e) {
      self.postMessage({ type: 'result', id: msg.id, error: String(e && e.message ? e.message : e) });
    }
  }
};