// Browser side of the end-to-end decision check: WebBrain's own prompt helpers (webgpu-decisions.js) + the MiniCPM-V
// runtime + the trained probe, on d1's balanced-V3 split screenshots, exactly as the extension would call it.
import * as ort from '/vendor/ort.webgpu.mjs';
import { AutoTokenizer, env } from '/vendor/transformers.web.js';
import { MiniCPMVRuntime } from '/runtime/minicpmv-runtime.mjs';
import { runMiniCPMVDecision } from '/runtime/minicpmv-decisions.mjs';

async function loadRGB(url) {
  const blob = await (await fetch(url)).blob();
  const bmp = await createImageBitmap(blob, { colorSpaceConversion: 'none', premultiplyAlpha: 'none' });
  const canvas = new OffscreenCanvas(bmp.width, bmp.height);
  const ctx = canvas.getContext('2d', { willReadFrequently: true });
  ctx.drawImage(bmp, 0, 0);
  const rgba = ctx.getImageData(0, 0, bmp.width, bmp.height).data;
  const rgb = new Uint8Array(bmp.width * bmp.height * 3);
  for (let i = 0, j = 0; i < rgba.length; i += 4, j += 3) { rgb[j] = rgba[i]; rgb[j + 1] = rgba[i + 1]; rgb[j + 2] = rgba[i + 2]; }
  return { rgb, width: bmp.width, height: bmp.height };
}

export async function run({ bundle = '/bundle', split = '/d1/d1-visual-adaptation-balanced-v3/data/final_test.json',
  requests = 1000, variant = 'probe', embedGraph = 'embed_tokens_q8f16.onnx', decoderGraph = 'decoder_model_merged_q8f16.onnx',
  visionGraph = 'vision_encoder_fp16.onnx', probeModule = '/runtime/minicpmv-decision-probe.js' } = {}) {
  env.allowLocalModels = true;
  env.allowRemoteModels = false;
  env.localModelPath = '/';
  ort.env.wasm.wasmPaths = '/vendor/';
  ort.env.wasm.numThreads = 1;
  const tokenizer = await AutoTokenizer.from_pretrained(bundle.replace(/^\//, ''));
  const tok = { encode: text => tokenizer.encode(text, { add_special_tokens: false }) };
  const t0 = performance.now();
  const rt = await MiniCPMVRuntime.create(ort, tok, {
    embed: `${bundle}/onnx/${embedGraph}`, vision: `${bundle}/onnx/${visionGraph}`, decoder: `${bundle}/onnx/${decoderGraph}` });
  const loadMs = performance.now() - t0;
  let probe = null;
  if (variant === 'probe') {
    const m = await import(probeModule);
    const fn = Object.values(m).find(v => typeof v === 'function');
    probe = { scores: fn };
  }
  const data = await (await fetch(split)).json();
  const answers = {}, timings = [];
  for (const req of data.requests.slice(0, requests)) {
    const images = await Promise.all((req.images || [req.image]).map(p => loadRGB(`/d1${p}`)));
    const entries = req.questions.map(e => [e.id, e.question]);
    const started = performance.now();
    const out = await runMiniCPMVDecision(rt, { state: req.state, images, entries, probe, variant, returnHidden: false });
    timings.push({ id: req.id, total_ms: Math.round(performance.now() - started), ...out.diagnostics });
    Object.assign(answers, out.answers);
    console.log('[decision]', req.id, JSON.stringify(out.diagnostics));
  }
  const mean = k => Math.round(timings.reduce((s, t) => s + (Array.isArray(t[k]) ? t[k].reduce((a, b) => a + b, 0) : t[k]), 0) / timings.length);
  return { variant, graphs: { embedGraph, decoderGraph, visionGraph }, load_ms: Math.round(loadMs), answers, timings,
    mean_ms: { total: mean('total_ms'), vision: mean('vision_ms'), prefill: mean('prefill_ms'), questions: mean('question_ms') } };
}