// The WhatDoWeSee page: pick a model, drop an image, show its tags. Runs entirely in the browser.
import { MODELS, fileUrl, run } from "./tagger.js";

/* global ort */
ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/";

const $ = (id) => document.getElementById(id);
const loaded = new Map(); // model id -> Promise<{session, config}>
let current = MODELS[0];
let lastImage = null; // the File on screen, re-tagged when the model changes

function load(model) {
  if (!loaded.has(model.id)) {
    loaded.set(
      model.id,
      (async () => {
        const config = await (await fetch(fileUrl(model, "config.json"))).json();
        const session = await ort.InferenceSession.create(fileUrl(model, "model.onnx"), {
          executionProviders: ["wasm"],
        });
        return { session, config };
      })(),
    );
  }
  return loaded.get(model.id);
}

/** The image as RGBA pixels squashed to size × size (no crop), like the training preprocessing. */
async function pixels(file, size) {
  let bitmap;
  try {
    bitmap = await createImageBitmap(file, { resizeWidth: size, resizeHeight: size, resizeQuality: "high" });
  } catch {
    bitmap = await createImageBitmap(file);
  }
  const canvas = new OffscreenCanvas(size, size);
  const context = canvas.getContext("2d");
  context.imageSmoothingEnabled = true;
  context.imageSmoothingQuality = "high";
  context.drawImage(bitmap, 0, 0, size, size);
  bitmap.close();
  return context.getImageData(0, 0, size, size).data;
}

function render(tags, ms) {
  const present = tags.filter((t) => t.present);
  const chips = $("present");
  chips.replaceChildren(
    ...(present.length
      ? present.map((t) => {
          const li = document.createElement("li");
          li.textContent = t.name;
          li.title = t.definition;
          return li;
        })
      : [Object.assign(document.createElement("li"), { className: "none", textContent: "No attribute found" })]),
  );
  $("rows").replaceChildren(
    ...tags.map((t) => {
      const tr = document.createElement("tr");
      if (t.present) tr.className = "on";
      const name = document.createElement("td");
      name.textContent = t.name;
      name.title = t.definition;
      const barCell = document.createElement("td");
      const track = document.createElement("div");
      track.className = "track";
      const fill = document.createElement("div");
      fill.className = "fill";
      fill.style.width = `${(t.p * 100).toFixed(1)}%`;
      const tick = document.createElement("div");
      tick.className = "tick";
      tick.style.left = `calc(${(t.threshold * 100).toFixed(1)}% - 1px)`;
      tick.title = `threshold ${(t.threshold * 100).toFixed(0)}%`;
      track.append(fill, tick);
      barCell.append(track);
      const pct = document.createElement("td");
      pct.className = "num";
      pct.textContent = `${Math.round(t.p * 100)}`;
      tr.append(name, barCell, pct);
      return tr;
    }),
  );
  $("all").hidden = false;
  $("status").textContent = `${present.length} of ${tags.length} attributes · ${ms} ms in your browser · ${current.name}`;
}

async function tag(file) {
  if (!file || !file.type.startsWith("image/")) {
    $("status").textContent = "That is not an image.";
    return;
  }
  lastImage = file;
  const preview = $("preview");
  if (preview.src) URL.revokeObjectURL(preview.src);
  preview.src = URL.createObjectURL(file);
  preview.hidden = false;
  $("drop").classList.add("has-image");
  $("hint-title").textContent = "Drop, click or paste";
  $("hint-sub").textContent = "another image to tag it";
  $("status").textContent = "Tagging…";
  try {
    const { session, config } = await load(current);
    const started = performance.now();
    const tags = await run(ort, session, config, await pixels(file, config.input.shape[3]));
    render(tags, Math.round(performance.now() - started));
  } catch (error) {
    $("status").textContent = `Could not tag this image: ${error.message ?? error}`;
  }
}

function setup() {
  const select = $("model");
  for (const m of MODELS) {
    const option = document.createElement("option");
    option.value = m.id;
    option.textContent = m.name;
    select.append(option);
  }
  select.addEventListener("change", () => {
    current = MODELS.find((m) => m.id === select.value);
    $("model-link").href = `https://huggingface.co/${current.repo}`;
    if (lastImage) tag(lastImage);
    else warm();
  });

  const drop = $("drop");
  const input = $("file");
  drop.addEventListener("click", () => input.click());
  drop.addEventListener("keydown", (e) => {
    if (e.key === "Enter" || e.key === " ") {
      e.preventDefault();
      input.click();
    }
  });
  input.addEventListener("change", () => tag(input.files[0]));
  for (const type of ["dragenter", "dragover"]) {
    drop.addEventListener(type, (e) => {
      e.preventDefault();
      drop.classList.add("over");
    });
  }
  for (const type of ["dragleave", "drop"]) {
    drop.addEventListener(type, () => drop.classList.remove("over"));
  }
  drop.addEventListener("drop", (e) => {
    e.preventDefault();
    tag(e.dataTransfer.files[0]);
  });
  document.addEventListener("paste", (e) => {
    const item = [...e.clipboardData.items].find((i) => i.type.startsWith("image/"));
    if (item) tag(item.getAsFile());
  });
  warm();
}

/** Load the model up front, so the first image is tagged right away. */
function warm() {
  $("status").textContent = "Loading the model (about 16 MB, once)…";
  load(current).then(
    () => {
      if (!lastImage) $("status").textContent = "Ready. Drop an image to tag it.";
    },
    (error) => {
      $("status").textContent = `Could not load the model: ${error.message ?? error}`;
    },
  );
}

setup();