// Browser side of the WebGPU check (loaded by harness.mjs into headless Chromium).
// 1. preprocessing parity: JS slicing/resize/normalize, prompt ids and canvas M-RoPE vs the Python processor;
// 2. model parity on WebGPU: verify margins (Yes - No) and decision_hidden vs PyTorch fp32, plus timings.
import * as ort from '/vendor/ort.webgpu.mjs';
import { AutoTokenizer, env } from '/vendor/transformers.web.js';
import { MINICPMV, MiniCPMVRuntime, canvasPositions, expandImagePlaceholders, preprocessImage, f16ToF32 } from '/runtime/minicpmv-runtime.mjs';

const log = (...a) => console.log('[check]', ...a);

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 };
}

function cosine(a, b) {
  let d = 0, na = 0, nb = 0;
  for (let i = 0; i < a.length; i++) { d += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i]; }
  return d / Math.sqrt(na * nb);
}

function f32b64(arr) {
  const bytes = new Uint8Array(arr.buffer, arr.byteOffset, arr.byteLength);
  let bin = '';
  for (let i = 0; i < bytes.length; i += 8192) bin += String.fromCharCode(...bytes.subarray(i, i + 8192));
  return btoa(bin);
}

function b64f32(text) {
  const bin = atob(text);
  const bytes = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
  return new Float32Array(bytes.buffer);
}

export async function run({ bundle = '/bundle', fixtures = '/fixtures', requests = 6, model = true, adapterInfo = true,
  embedGraph = 'embed_tokens_q8f16.onnx', decoderGraph = 'decoder_model_merged_q8f16.onnx',
  visionGraph = 'vision_encoder_fp16.onnx', generateTest = true } = {}) {
  const result = { checks: [], preprocessing: [], decisions: [], timings: {}, errors: [] };
  if (adapterInfo && navigator.gpu) {
    const adapter = await navigator.gpu.requestAdapter({ powerPreference: 'high-performance' });
    result.adapter = adapter ? { ...adapter.info, vendor: adapter.info?.vendor, architecture: adapter.info?.architecture,
      description: adapter.info?.description, features: [...adapter.features].filter(f => f.includes('f16')) } : null;
  }
  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 fx = await (await fetch(`${fixtures}/fixtures.json`)).json();
  const reqs = fx.requests.slice(0, requests);

  // ---------------------------------------------------------------- preprocessing parity
  const images = [];
  for (const r of reqs) {
    const t0 = performance.now();
    const { rgb, width, height } = await loadRGB(`/d1${r.image}`);
    const image = preprocessImage(rgb, width, height);
    const prepMs = performance.now() - t0;
    images.push(image);
    const ref = r.crops_bin ? new Float32Array(await (await fetch(`${fixtures}/${r.crops_bin}`)).arrayBuffer()) : null;
    let off = 0, maxDiff = ref ? 0 : null, nDiff = ref ? 0 : null;
    const sizesOk = JSON.stringify(image.crops.map(c => [c.h, c.w])) === JSON.stringify(r.target_sizes);
    for (const c of image.crops) {
      if (!ref) break;
      for (let i = 0; i < c.pixels.length; i++) {
        const d = Math.abs(c.pixels[i] - ref[off + i]);
        if (d > 0) nDiff++;
        if (d > maxDiff) maxDiff = d;
      }
      off += c.pixels.length;
    }
    const text = expandImagePlaceholders(r.prefix, [image]);
    const ids = tok.encode(text);
    const idsEqual = ids.length === r.input_ids.length && ids.every((v, i) => v === r.input_ids[i]);
    const pos = canvasPositions(ids, image.crops.map(c => [c.h, c.w]));
    let posMismatch = 0;
    for (let i = 0; i < ids.length && idsEqual; i++) {
      if (pos.t[i] !== r.positions[0][i] || pos.h[i] !== r.positions[1][i] || pos.w[i] !== r.positions[2][i]) posMismatch++;
    }
    const row = { id: r.id, crops: image.crops.length, crop_sizes_equal: sizesOk, pixel_max_abs_diff: maxDiff,
      pixel_values_differing: nDiff, pixel_values: ref ? ref.length : null, ids_equal: idsEqual, tokens: ids.length,
      position_mismatches: posMismatch, preprocess_ms: Math.round(prepMs) };
    result.preprocessing.push(row);
    log(JSON.stringify(row));
  }
  if (!model) return result;

  // ---------------------------------------------------------------- WebGPU model parity
  const t0 = performance.now();
  const rt = await MiniCPMVRuntime.create(ort, tok, {
    embed: `${bundle}/onnx/${embedGraph}`,
    vision: `${bundle}/onnx/${visionGraph}`,
    decoder: `${bundle}/onnx/${decoderGraph}`,
  });
  result.graphs = { embedGraph, visionGraph, decoderGraph };
  result.timings.load_ms = Math.round(performance.now() - t0);
  const yes = fx.yes_id, no = fx.no_id;
  const prefillMs = [], questionMs = [], visionMs = [];
  const answers = {};
  result.report = { engines: { webgpu: { answers } } };
  for (let ri = 0; ri < reqs.length; ri++) {
    const r = reqs[ri], image = images[ri];
    const tv = performance.now();
    const inputs = await rt.buildInputs(r.prefix, [image]);
    visionMs.push(performance.now() - tv);
    const tp = performance.now();
    const pre = await rt.forward(inputs.embeds, inputs.pos, 0, inputs.ids.length, rt.emptyCache());
    prefillMs.push(performance.now() - tp);
    const P = inputs.ids.length;
    for (const q of r.questions) {
      const tq = performance.now();
      const margins = [], cos = [], hiddens = [];
      for (let oi = 0; oi < q.suffix_ids.length; oi++) {
        const sids = q.suffix_ids[oi];
        const full = inputs.ids.concat(sids);
        const pos = canvasPositions(full, inputs.sizes);
        const emb = await rt.embedIds(sids);
        const shifted = { t: pos.t.subarray(P), h: pos.h.subarray(P), w: pos.w.subarray(P) };
        const out = await rt.forward(emb, shifted, 0, sids.length, pre.cache, { keepCache: false });
        margins.push(out.logits[yes] - out.logits[no]);
        hiddens.push(f32b64(out.hidden));
        if (q.hidden_b64) cos.push(cosine(out.hidden, b64f32(q.hidden_b64[oi])));
      }
      questionMs.push(performance.now() - tq);
      const argmax = a => a.indexOf(Math.max(...a));
      const row = { id: q.id, type: q.question.type, webgpu_margins: margins.map(v => +v.toFixed(3)),
        pt_margins: q.margins?.map(v => +v.toFixed(3)), argmax_equal: q.margins ? argmax(margins) === argmax(q.margins) : null,
        max_margin_diff: q.margins ? Math.max(...margins.map((v, i) => Math.abs(v - q.margins[i]))) : null,
        hidden_cos_min: cos.length ? Math.min(...cos) : null };
      result.decisions.push(row);
      answers[q.id] = { margins, hidden_f32: hiddens };
      log(JSON.stringify(row));
    }
    MiniCPMVRuntime.disposeCache(pre.cache);
  }
  // generation smoke test: one screenshot, one free-text question
  if (generateTest) {
    const r = reqs[0];
    const prompt = `<|im_start|>user\n${MINICPMV.imagePad}\nDescribe this web page in one sentence.<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n`;
    const tg = performance.now();
    const gen = await rt.generate(prompt, [images[0]], { maxNewTokens: 48 });
    const ms = performance.now() - tg;
    result.generation = { request: r.id, prompt_tokens: gen.promptTokens, new_tokens: gen.ids.length,
      text: tokenizer.decode(gen.ids, { skip_special_tokens: true }), ms: Math.round(ms),
      decode_tokens_per_s: null };
    log(JSON.stringify(result.generation));
  }
  const mean = a => a.reduce((s, v) => s + v, 0) / Math.max(1, a.length);
  result.timings = { ...result.timings, vision_and_embed_ms_mean: Math.round(mean(visionMs)), prefill_ms_mean: Math.round(mean(prefillMs)),
    question_ms_mean: Math.round(mean(questionMs)) };
  const d = result.decisions.filter(x => x.argmax_equal !== null);
  result.summary = { questions: d.length, argmax_agreement: d.filter(x => x.argmax_equal).length / Math.max(1, d.length),
    worst_margin_diff: Math.max(...d.map(x => x.max_margin_diff)), mean_margin_diff: mean(d.map(x => x.max_margin_diff)),
    min_hidden_cos: Math.min(...d.map(x => x.hidden_cos_min ?? 1)) };
  return result;
}