// Image validation and model preprocessing, ported from src/artifex/preprocess.py.
//
// The resize is a line-by-line port of Pillow's fixed-point bicubic resampler (libImaging/Resample.c,
// 8 bits per channel), so the 224 x 224 model input matches the Python pipeline pixel for pixel once
// the image is decoded. Decoders themselves (browser vs. Pillow/libjpeg) can differ slightly; the
// measured end-to-end effect is reported on the Methodology page.

export class ImageValidationError extends Error {
  constructor(code, message) {
    super(message);
    this.code = code;
  }
}

export function sniffFormat(b) {
  if (b.length >= 12 && b[0] === 0x52 && b[1] === 0x49 && b[2] === 0x46 && b[3] === 0x46 && b[8] === 0x57 && b[9] === 0x45 && b[10] === 0x42 && b[11] === 0x50)
    return "WEBP";
  if (b.length >= 3 && b[0] === 0xff && b[1] === 0xd8 && b[2] === 0xff) return "JPEG";
  const png = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
  if (b.length >= 8 && png.every((v, i) => b[i] === v)) return "PNG";
  return null;
}

// Width and height from the file header (before decoding), or null if they cannot be read.
export function headerSize(b, fmt) {
  const u16be = (o) => (b[o] << 8) | b[o + 1];
  const u32be = (o) => ((b[o] << 24) >>> 0) + (b[o + 1] << 16) + (b[o + 2] << 8) + b[o + 3];
  const u24le = (o) => b[o] | (b[o + 1] << 8) | (b[o + 2] << 16);
  try {
    if (fmt === "PNG") return b.length >= 24 ? [u32be(16), u32be(20)] : null;
    if (fmt === "JPEG") {
      let o = 2;
      while (o + 9 < b.length) {
        if (b[o] !== 0xff) return null;
        let m = b[o + 1];
        if (m === 0xff) {
          o += 1;
          continue;
        }
        if (m === 0xd8 || m === 0x01 || (m >= 0xd0 && m <= 0xd7)) {
          o += 2;
          continue;
        }
        const len = u16be(o + 2);
        if (m >= 0xc0 && m <= 0xcf && m !== 0xc4 && m !== 0xc8 && m !== 0xcc) return [u16be(o + 7), u16be(o + 5)];
        o += 2 + len;
      }
      return null;
    }
    if (fmt === "WEBP") {
      const tag = String.fromCharCode(b[12], b[13], b[14], b[15]);
      if (tag === "VP8 ") return [(b[26] | (b[27] << 8)) & 0x3fff, (b[28] | (b[29] << 8)) & 0x3fff];
      if (tag === "VP8L") {
        const bits = b[21] | (b[22] << 8) | (b[23] << 16) | (b[24] << 24);
        return [(bits & 0x3fff) + 1, ((bits >> 14) & 0x3fff) + 1];
      }
      if (tag === "VP8X") return [u24le(24) + 1, u24le(27) + 1];
    }
  } catch {
    return null;
  }
  return null;
}

// Decode and validate untrusted bytes. Returns {rgb: Uint8Array (w*h*3), width, height}.
export async function loadImageBytes(bytes, { allowedFormats, maxBytes, maxPixels, minSide }) {
  if (!bytes.length) throw new ImageValidationError("empty", "The file is empty.");
  if (bytes.length > maxBytes) throw new ImageValidationError("too-large", `The file is larger than the ${Math.floor(maxBytes / 1e6)} MB limit.`);
  const fmt = sniffFormat(bytes);
  if (!fmt || !allowedFormats.includes(fmt))
    throw new ImageValidationError("unsupported-format", "Unsupported file type. Please upload a JPEG, PNG or WebP image.");
  const dims = headerSize(bytes, fmt);
  if (dims && dims[0] && dims[1] && dims[0] * dims[1] > maxPixels)
    throw new ImageValidationError("too-many-pixels", `The image is too large to analyse (${dims[0]}×${dims[1]} pixels).`);

  let bitmap;
  try {
    const type = { JPEG: "image/jpeg", PNG: "image/png", WEBP: "image/webp" }[fmt];
    bitmap = await createImageBitmap(new Blob([bytes], { type }), {
      imageOrientation: "from-image",
      premultiplyAlpha: "none",
      colorSpaceConversion: "none",
    });
  } catch (e) {
    throw new ImageValidationError("corrupt", `The image could not be decoded (${e?.name || "DecodeError"}).`);
  }
  const width = bitmap.width;
  const height = bitmap.height;
  if (width * height > maxPixels) {
    bitmap.close();
    throw new ImageValidationError("too-many-pixels", `The image is too large to analyse (${width}×${height} pixels).`);
  }
  if (Math.min(width, height) < minSide) {
    bitmap.close();
    throw new ImageValidationError("too-small", `The image is too small (minimum side ${minSide} pixels).`);
  }
  let rgba;
  try {
    const canvas = typeof OffscreenCanvas !== "undefined" ? new OffscreenCanvas(width, height) : Object.assign(document.createElement("canvas"), { width, height });
    const ctx = canvas.getContext("2d", { willReadFrequently: true, colorSpace: "srgb" });
    ctx.drawImage(bitmap, 0, 0);
    rgba = ctx.getImageData(0, 0, width, height, { colorSpace: "srgb" }).data;
  } catch {
    throw new ImageValidationError("too-many-pixels", `This browser cannot process an image of ${width}×${height} pixels. Please use a smaller copy.`);
  } finally {
    bitmap.close();
  }
  return { rgb: toRgb(rgba, width * height), width, height };
}

// RGBA -> RGB, compositing transparency onto white exactly as Pillow's paste() with an alpha mask does.
export function toRgb(rgba, n) {
  const out = new Uint8Array(n * 3);
  for (let i = 0; i < n; i++) {
    const a = rgba[i * 4 + 3];
    for (let c = 0; c < 3; c++) {
      const s = rgba[i * 4 + c];
      if (a === 255) out[i * 3 + c] = s;
      else {
        const t = s * a + 255 * (255 - a) + 128; // BLEND(mask, in1=white, in2=src) with DIV255
        out[i * 3 + c] = ((t >> 8) + t) >> 8;
      }
    }
  }
  return out;
}

// ---------------------------------------------------------------- Pillow bicubic resampling (8 bpc)

const PRECISION_BITS = 32 - 8 - 2;
const ONE = 2 ** PRECISION_BITS;
const HALF = 2 ** (PRECISION_BITS - 1);

function bicubic(x) {
  const a = -0.5;
  if (x < 0) x = -x;
  if (x < 1) return ((a + 2) * x - (a + 3)) * x * x + 1;
  if (x < 2) return (((x - 5) * x + 8) * x - 4) * a;
  return 0;
}

function precomputeCoeffs(inSize, outSize) {
  const scale = inSize / outSize;
  const filterscale = scale < 1 ? 1 : scale;
  const support = 2.0 * filterscale;
  const ksize = Math.ceil(support) * 2 + 1;
  const kk = new Int32Array(outSize * ksize);
  const bounds = new Int32Array(outSize * 2);
  const pre = new Float64Array(ksize);
  for (let xx = 0; xx < outSize; xx++) {
    const center = (xx + 0.5) * scale;
    const ss = 1.0 / filterscale;
    let ww = 0;
    let xmin = Math.trunc(center - support + 0.5);
    if (xmin < 0) xmin = 0;
    let xmax = Math.trunc(center + support + 0.5);
    if (xmax > inSize) xmax = inSize;
    xmax -= xmin;
    pre.fill(0);
    for (let x = 0; x < xmax; x++) {
      const w = bicubic((x + xmin - center + 0.5) * ss);
      pre[x] = w;
      ww += w;
    }
    for (let x = 0; x < xmax; x++) if (ww !== 0) pre[x] /= ww;
    for (let x = 0; x < ksize; x++) {
      const v = pre[x];
      kk[xx * ksize + x] = v < 0 ? Math.trunc(-0.5 + v * ONE) : Math.trunc(0.5 + v * ONE);
    }
    bounds[xx * 2] = xmin;
    bounds[xx * 2 + 1] = xmax;
  }
  return { ksize, kk, bounds };
}

const clip8 = (ss) => {
  const v = Math.floor(ss / ONE); // arithmetic shift right by PRECISION_BITS
  return v < 0 ? 0 : v > 255 ? 255 : v;
};

// Resize an RGB image (Uint8Array, w*h*3) to (nw, nh) like PIL.Image.resize(..., Image.BICUBIC).
export function resizeBicubic(src, w, h, nw, nh) {
  const needH = nw !== w;
  const needV = nh !== h;
  if (!needH && !needV) return src.slice();
  const horiz = precomputeCoeffs(w, nw);
  const vert = precomputeCoeffs(h, nh);
  let img = src;
  let iw = w;
  const vb = vert.bounds.slice();
  if (needH) {
    const yFirst = vert.bounds[0];
    const yLast = vert.bounds[(nh - 1) * 2] + vert.bounds[(nh - 1) * 2 + 1];
    if (needV) for (let i = 0; i < nh; i++) vb[i * 2] -= yFirst;
    const rows = needV ? yLast - yFirst : h;
    const off = needV ? yFirst : 0;
    const tmp = new Uint8Array(nw * rows * 3);
    const { ksize, kk, bounds } = horiz;
    for (let yy = 0; yy < rows; yy++) {
      const rowIn = (yy + off) * w * 3;
      for (let xx = 0; xx < nw; xx++) {
        const xmin = bounds[xx * 2];
        const xmax = bounds[xx * 2 + 1];
        const k = xx * ksize;
        let s0 = HALF;
        let s1 = HALF;
        let s2 = HALF;
        for (let x = 0; x < xmax; x++) {
          const p = rowIn + (x + xmin) * 3;
          const c = kk[k + x];
          s0 += src[p] * c;
          s1 += src[p + 1] * c;
          s2 += src[p + 2] * c;
        }
        const o = (yy * nw + xx) * 3;
        tmp[o] = clip8(s0);
        tmp[o + 1] = clip8(s1);
        tmp[o + 2] = clip8(s2);
      }
    }
    img = tmp;
    iw = nw;
  }
  if (!needV) return img;
  const out = new Uint8Array(nw * nh * 3);
  const { ksize, kk } = vert;
  for (let yy = 0; yy < nh; yy++) {
    const ymin = vb[yy * 2];
    const ymax = vb[yy * 2 + 1];
    const k = yy * ksize;
    for (let xx = 0; xx < iw; xx++) {
      let s0 = HALF;
      let s1 = HALF;
      let s2 = HALF;
      for (let y = 0; y < ymax; y++) {
        const p = ((y + ymin) * iw + xx) * 3;
        const c = kk[k + y];
        s0 += img[p] * c;
        s1 += img[p + 1] * c;
        s2 += img[p + 2] * c;
      }
      const o = (yy * nw + xx) * 3;
      out[o] = clip8(s0);
      out[o + 1] = clip8(s1);
      out[o + 2] = clip8(s2);
    }
  }
  return out;
}

// Python's round(): halves go to the even neighbour.
export function pyRound(x) {
  const f = Math.floor(x);
  const d = x - f;
  if (d > 0.5) return f + 1;
  if (d < 0.5) return f;
  return f % 2 === 0 ? f : f + 1;
}

// Letterbox to size x size on an ImageNet-mean background (geometric_transform, geometry="letterbox").
export function letterbox(img, spec) {
  const s = spec.image_size;
  const { width: w, height: h } = img;
  const scale = s / Math.max(w, h);
  const nw = Math.max(1, pyRound(w * scale));
  const nh = Math.max(1, pyRound(h * scale));
  const resized = resizeBicubic(img.rgb, w, h, nw, nh);
  const fill = spec.mean.map((m) => pyRound(m * 255));
  const canvas = new Uint8Array(s * s * 3);
  for (let i = 0; i < s * s; i++) {
    canvas[i * 3] = fill[0];
    canvas[i * 3 + 1] = fill[1];
    canvas[i * 3 + 2] = fill[2];
  }
  const ox = Math.floor((s - nw) / 2);
  const oy = Math.floor((s - nh) / 2);
  for (let y = 0; y < nh; y++) canvas.set(resized.subarray(y * nw * 3, (y + 1) * nw * 3), ((y + oy) * s + ox) * 3);
  return canvas;
}

// float32 CHW tensor normalised with the encoder's mean/std (to_model_array), float32 rounding at each step.
export function toModelArray(img, spec) {
  const s = spec.image_size;
  const px = letterbox(img, spec);
  const out = new Float32Array(3 * s * s);
  const f = Math.fround;
  const mean = spec.mean.map(f);
  const std = spec.std.map(f);
  for (let i = 0; i < s * s; i++) {
    for (let c = 0; c < 3; c++) out[c * s * s + i] = f(f(f(px[i * 3 + c] / 255) - mean[c]) / std[c]);
  }
  return out;
}