import type * as Ort from 'onnxruntime-web';

// Inventing pixels, in the browser. A small Real-ESRGAN — realesr-general-x4v3,
// a 4.9MB SRVGGNetCompact that turns one pixel into four — runs through
// onnxruntime-web on the visitor's own machine. The photo never leaves the
// device, so the API's promise ("Never sees a photo") survives an export that
// asks for more pixels than the picture has.
//
// Both the runtime (served next to canvaskit.wasm, see scripts/copy-wasm.mjs)
// and the model (public/models) are static files; nothing here changes what the
// server can do.
//
// The model is not the file as published: its 33 PReLU activations are written
// out as Relu/Neg/Mul/Sub, because the WebGPU provider has no PReLU kernel and
// would otherwise hand every activation to the processor and take it back, 33
// times per tile. scripts/realesr-gpu.py is that rewrite and explains it.

const MODEL_URL = '/models/realesr-general-x4v3.onnx';
// The same graph in half precision: 2.3MB of weights against 4.9, and twice the
// arithmetic per pass on a chip whose shaders multiply in fp16. Which copy is
// asked for is decided at load time, never here — see `openSession`.
const MODEL_F16_URL = '/models/realesr-general-x4v3-f16.onnx';
// The wasm dir, not the file: the runtime picks its own name inside it.
const WASM_DIR = '/wasm/ort/';
// What the export menu warms, named so the warming costs no runtime: the script
// the runtime hands its threads, the binary they all load, and the two copies of
// the model. See `preloadSuperRes`.
const RUNTIME_FILES = ['ort-wasm-simd-threaded.jsep.mjs', 'ort-wasm-simd-threaded.jsep.wasm'];
// The model's own factor is fixed at 4; the destination is whatever the export
// asked for, so each finished tile is drawn to its place in the output rather
// than assembled at 4x and shrunk (which would cost the memory of both).
const MODEL_SCALE = 4;
// Source pixels per pass, and the margin the model sees but the draw discards —
// without it every tile boundary shows as a seam. Counted in whatever the model
// is being fed, which is not always the photo's own pixels: see `feedScale`.
const TILE = 256;
const PAD = 8;

export interface UpscaleProgress {
  done: number;
  total: number;
}

interface Loaded {
  ort: typeof Ort;
  session: Ort.InferenceSession;
}

let loaded: Promise<Loaded> | null = null;

// Does this browser's GPU multiply in half precision? Not "does it have WebGPU":
// the fp16 weights are only worth asking for where the adapter says `shader-f16`,
// because the provider emulates the type without it — same file, same maths, no
// speed. An adapter too is a promise the chip may not have.
async function canUseF16(): Promise<boolean> {
  const gpu = (navigator as unknown as { gpu?: { requestAdapter(o?: unknown): Promise<{ features: Set<string> } | null> } }).gpu;
  if (!gpu) return false;
  try {
    const adapter = await gpu.requestAdapter({ powerPreference: 'high-performance' });
    return adapter?.features?.has('shader-f16') === true;
  } catch {
    return false;
  }
}

// WebGPU first — an export waits seconds instead of a minute — and the same jsep
// build falls back to plain WebAssembly when the browser has no GPU adapter to
// give. Three attempts at most, in that order: the half-precision model on the
// GPU, the full-precision one there, and the full-precision one on the
// processor, which is the only copy the WebAssembly provider can run.
async function openSession(ort: typeof Ort): Promise<Ort.InferenceSession> {
  const onGpu = (url: string) =>
    ort.InferenceSession.create(url, { executionProviders: ['webgpu'] }).catch(() => null);
  if (await canUseF16()) {
    const half = await onGpu(MODEL_F16_URL);
    if (half) return half;
  }
  return (await onGpu(MODEL_URL)) ?? ort.InferenceSession.create(MODEL_URL, { executionProviders: ['wasm'] });
}

// Loads the runtime and the model once per page. See `openSession` for which
// model, and `realesr-fp16.py` for what the half-precision copy is.
function load(): Promise<Loaded> {
  if (!loaded) {
    loaded = (async () => {
      const ort = await import('onnxruntime-web');
      ort.env.wasm.wasmPaths = WASM_DIR;
      // nginx sends COOP+COEP, so the page is cross-origin isolated and there is
      // a SharedArrayBuffer for the threaded build to spread a tile over. The
      // check is for the case the headers ever go missing: asking for more than
      // one thread without isolation makes the runtime throw instead of falling
      // back, and a slow export beats a broken one.
      //
      // Four, not the processor's whole width. Every thread is a Web Worker of
      // its own holding a copy of that 27MB binary, and on a laptop the memory
      // and the bandwidth are both spent well before the eighth one: the tile is
      // 256 pixels square, and the threads that matter are here. A machine with
      // fewer cores asks for fewer.
      ort.env.wasm.numThreads = self.crossOriginIsolated
        ? Math.min(4, navigator.hardwareConcurrency || 1)
        : 1;
      // A laptop with two GPUs would otherwise hand this to the one built into
      // the processor. The model is the export, so ask for the fast one.
      ort.env.webgpu.powerPreference = 'high-performance';
      // env.logLevel looks like a switch for the runtime's own warnings and is not
      // one — it is read when the module initialises, before this runs, so a level
      // set here silences nothing. Left alone, they are worth reading: the last one
      // to appear said the graph was being split between the chip and the processor.
      const session = await openSession(ort);
      // One throwaway run through the graph: the runtime builds and caches its
      // kernels on the first one, and that is a wait the visitor should spend
      // looking at the export menu rather than at a progress bar. Sixteen
      // pixels: inaudible next to the load, and it keeps the kernels out of the
      // way of the tiles that follow.
      await session
        .run({ [session.inputNames[0]]: new ort.Tensor('float32', new Float32Array(3 * 16 * 16), [1, 3, 16, 16]) })
        .catch(() => {});
      return { ort, session };
    })().catch((err) => {
      // A failed load must not poison the next attempt.
      loaded = null;
      throw err;
    });
  }
  return loaded;
}

// Start fetching the 32MB of runtime and model before the export that needs
// them, so the visitor waits for the pixels and not for the download. Called
// when the export menu opens — the only door onto an upscale.
//
// The bytes, and only the bytes: they are fetched by name rather than imported,
// because importing the runtime is what starts its thread pool — one Web Worker
// per thread, one copy of the 27MB binary each — and a pool that then belongs to
// the page until it dies, whether or not an export is ever asked for. Opening a
// menu is not the work. What the menu is worth is the download, and the download
// is what the browser's cache keeps: the export's own `import` then pays a second
// of instantiation it cannot avoid anyway, and finds the files already on disk.
export function preloadSuperRes(): void {
  for (const file of [...RUNTIME_FILES.map((f) => WASM_DIR + f), MODEL_URL, MODEL_F16_URL]) {
    void fetch(file).catch(() => {});
  }
}

// The page is being put away — closed, reloaded, restored from the back/forward
// cache — and this renderer is holding the session's own buffers and a thread
// pool with it. Hand the session back rather than leaving the browser to reap it
// with the page: a machine short of memory at the moment the studio closes is a
// machine that hangs on the studio opening again, which is the one symptom this
// can answer for.
//
// The pool is the runtime's and goes with the page; the session is ours. `loaded`
// goes with it, so a page that comes back loads the runtime again instead of
// calling into a released session — the files are cached, so that is a second.
export async function releaseSuperRes(): Promise<void> {
  const pending = loaded;
  loaded = null;
  if (!pending) return;
  try {
    await (await pending).session.release();
  } catch {
    // A load that failed, or a runtime too old to have `release`: nothing to
    // give back, and nothing to say about it either.
  }
}

// The last moment the studio is on screen. Nothing else in this module listens
// for anything, so the listener goes with the one resource that needs it.
if (typeof addEventListener === 'function') {
  addEventListener('pagehide', () => {
    void releaseSuperRes();
  });
}

// The photo's bytes, enlarged so its longest edge is `targetLongest`. The caller
// only gets here when the target is larger than the picture (App decides), and
// gets back the same bytes untouched when it is not.
export async function upscaleJpeg(
  bytes: Uint8Array,
  targetLongest: number,
  onProgress?: (p: UpscaleProgress) => void
): Promise<Uint8Array> {
  const bitmap = await createImageBitmap(new Blob([bytes as BlobPart]));
  try {
    const w = bitmap.width;
    const h = bitmap.height;
    const longest = Math.max(w, h);
    if (targetLongest <= longest) return bytes;
    const scale = targetLongest / longest;

    const dstCanvas = new OffscreenCanvas(Math.max(1, Math.round(w * scale)), Math.max(1, Math.round(h * scale)));
    // Opaque: a partly covered edge pixel would otherwise survive as transparency
    // and the JPEG export flattens that onto black — a dark line down every seam.
    const dstCtx = dstCanvas.getContext('2d', { alpha: false });
    if (!dstCtx) return bytes;
    // The destination is rarely the model's own 4x, so the draw below is usually
    // a reduction, by up to four, and `low` would keep one sample in four of what
    // the model has just drawn. `high` reads them all.
    dstCtx.imageSmoothingEnabled = true;
    dstCtx.imageSmoothingQuality = 'high';

    // The photo's own pixels, read once. The model is handed the picture itself
    // and never a smaller copy of it: it answers with four pixels for every one
    // it is given, and a source shrunk towards the destination is detail the
    // photo had that the model is then asked to invent back — a 2px stripe in a
    // 2400px photo exported at 4K comes back as flat grey that way. The four for
    // one is spent in the draw below instead, which reduces to the destination
    // and keeps what the photo actually held.
    const feedCanvas = new OffscreenCanvas(w, h);
    const feedCtx = feedCanvas.getContext('2d', { willReadFrequently: true });
    if (!feedCtx) return bytes;
    feedCtx.drawImage(bitmap, 0, 0);
    const src = feedCtx.getImageData(0, 0, w, h);

    const { ort, session } = await load();
    const inputName = session.inputNames[0];
    const cols = Math.ceil(w / TILE);
    const rows = Math.ceil(h / TILE);
    // Destination pixels per fed pixel. Derived from the destination itself so
    // the last row and column land exactly on its edge rather than a rounding
    // short of it, and shared by neighbouring tiles so their boundary is the
    // same number for both and nothing is left half-covered.
    const stepX = dstCanvas.width / w;
    const stepY = dstCanvas.height / h;
    let done = 0;
    onProgress?.({ done, total: cols * rows });

    for (let ty = 0; ty < rows; ty++) {
      for (let tx = 0; tx < cols; tx++) {
        const x0 = tx * TILE;
        const y0 = ty * TILE;
        const tw = Math.min(TILE, w - x0);
        const th = Math.min(TILE, h - y0);
        // The margin the model gets: full on the inside, clipped at the photo's
        // own edge, so the tensor covers whole pixels only.
        const left = Math.min(PAD, x0);
        const top = Math.min(PAD, y0);
        const pw = tw + left + Math.min(PAD, w - (x0 + tw));
        const ph = th + top + Math.min(PAD, h - (y0 + th));

        // NCHW, 0..1 RGB — what the model was trained to read.
        const input = new Float32Array(3 * pw * ph);
        const plane = pw * ph;
        for (let y = 0; y < ph; y++) {
          const srow = ((y0 - top + y) * w + (x0 - left)) * 4;
          for (let x = 0; x < pw; x++) {
            const s = srow + x * 4;
            input[y * pw + x] = src.data[s] / 255;
            input[plane + y * pw + x] = src.data[s + 1] / 255;
            input[2 * plane + y * pw + x] = src.data[s + 2] / 255;
          }
        }

        const out = await session.run({ [inputName]: new ort.Tensor('float32', input, [1, 3, ph, pw]) });
        const tensor = out[session.outputNames[0]];
        const dims = tensor.dims as number[];
        const ow = dims[3];
        const oh = dims[2];
        const opts = tensor.data as Float32Array;

        // Paint the tile's own pixels into a canvas of their own — the padded ring
        // the model needed to see is skipped here instead of drawn and clipped —
        // then draw that canvas into the output at the output's scale.
        const cw = tw * MODEL_SCALE;
        const ch = th * MODEL_SCALE;
        const tile = new OffscreenCanvas(cw, ch);
        const tileCtx = tile.getContext('2d');
        if (!tileCtx) return bytes;
        const img = tileCtx.createImageData(cw, ch);
        const oplane = ow * oh;
        const ox = left * MODEL_SCALE;
        const oy = top * MODEL_SCALE;
        for (let y = 0; y < ch; y++) {
          const srcRow = (oy + y) * ow + ox;
          let d = y * cw * 4;
          for (let x = 0; x < cw; x++) {
            const s = srcRow + x;
            // A Uint8ClampedArray clamps and rounds on assignment; a clamp call
            // per channel was most of this loop's cost.
            img.data[d] = opts[s] * 255;
            img.data[d + 1] = opts[oplane + s] * 255;
            img.data[d + 2] = opts[2 * oplane + s] * 255;
            img.data[d + 3] = 255;
            d += 4;
          }
        }
        tileCtx.putImageData(img, 0, 0);
        // Destination edges snapped to whole pixels: neighbouring tiles then share
        // the exact same boundary, so nothing is left half-covered (and half-black)
        // where they meet. `stepX` is rarely an integer, so `x0 * stepX` alone
        // lands on a fraction and antialiases the seam.
        const dx0 = Math.round(x0 * stepX);
        const dy0 = Math.round(y0 * stepY);
        dstCtx.drawImage(
          tile,
          0,
          0,
          cw,
          ch,
          dx0,
          dy0,
          Math.round((x0 + tw) * stepX) - dx0,
          Math.round((y0 + th) * stepY) - dy0
        );

        done++;
        onProgress?.({ done, total: cols * rows });
      }
    }

    const blob = await dstCanvas.convertToBlob({ type: 'image/jpeg', quality: 0.95 });
    return new Uint8Array(await blob.arrayBuffer());
  } finally {
    bitmap.close();
  }
}