import sharp from "sharp";
import * as ort from "onnxruntime-web";
import os from "node:os";

function toCHW320(rgb: Buffer, size = 320) {
  const chw = new Float32Array(1 * 3 * size * size);
  let p = 0;
  for (let y = 0; y < size; y++) {
    for (let x = 0; x < size; x++) {
      const i = (y * size + x) * 3;
      const r = rgb[i] / 255; const g = rgb[i + 1] / 255; const b = rgb[i + 2] / 255;
      chw[p] = r; chw[p + size * size] = g; chw[p + 2 * size * size] = b; p++;
    }
  }
  return chw;
}

export async function inferU2MaskWasm(originalJpegOrPng: Buffer, modelUrl: string, wasmBaseUrl?: string): Promise<Uint8Array> {
  const wasmEnv = (ort.env as typeof ort.env & { wasm?: { wasmPaths?: unknown; simd?: boolean; numThreads?: number } }).wasm;
  if (wasmEnv) {
    if (wasmBaseUrl) {
      wasmEnv.wasmPaths = wasmBaseUrl as typeof wasmEnv.wasmPaths;
    }
    wasmEnv.simd = true;
    wasmEnv.numThreads = Math.max(1, Math.min(4, os.cpus()?.length ?? 1));
  }

  const img = sharp(originalJpegOrPng).removeAlpha();
  const meta = await img.metadata();
  const W = meta.width ?? 0, H = meta.height ?? 0;
  if (!(W > 0 && H > 0)) throw new Error("Bild hat keine gültigen Dimensionen");

  const target = 320;
  const resized = await img.resize(target, target, { fit: "fill" }).raw().toBuffer();
  const chw = toCHW320(resized, target);

  const session = await ort.InferenceSession.create(modelUrl, { executionProviders: ["wasm"] as any });
  const inputName = session.inputNames[0];
  const input = new ort.Tensor("float32", chw, [1, 3, target, target]);
  const results = await session.run({ [inputName]: input } as any);
  const firstOut = (results as any)[Object.keys(results)[0]] as ort.Tensor;
  const outData = firstOut.data as Float32Array;

  const mask320 = new Uint8Array(target * target);
  const n = Math.min(mask320.length, outData.length);
  for (let i = 0; i < n; i++) {
    let v = outData[i]; if (!Number.isFinite(v)) v = 0;
    v = 1 / (1 + Math.exp(-v));
    mask320[i] = Math.max(0, Math.min(255, Math.round(v * 255)));
  }

  const up = await sharp(Buffer.from(mask320), { raw: { width: target, height: target, channels: 1 } })
    .resize(W, H, { kernel: "lanczos3" })
    .toColorspace("b-w")
    .raw()
    .toBuffer();

  return new Uint8Array(up);
}