// Photo Enhancer (MegaPotato): enhance at the original size or 2×, compare two frames side by side.

import { DemoError, Demo, INFO_KINDS, NBSP, dims, errorCopy, escapeHtml, isResumable, thumbnail, usd } from "./core.js";

const SPACE_URL = "https://huggingface.co/spaces/MegaPotatoAI/Photo-Enhancer";
const COPY = { verb: "enhance", runs: "runs", spaceUrl: SPACE_URL };

const $ = (sel) => document.querySelector(sel);
const els = {
  notice: $("#notice"),
  file: $("#file"),
  thumb: $("#drop-thumb"),
  fileLabel: $("#file-label"),
  fileMeta: $("#file-meta"),
  plan: $("#plan"),
  quota: $("#quota"),
  run: $("#run"),
  status: $("#status"),
  caption: $("#caption"),
  examples: $("#examples"),
  credits: $("#credits"),
  beforeLabel: $("#before-label"),
  emptyNote: $(".pane-empty"),
  viewerWrap: $("#viewer-wrap"),
  inspect: $("#inspect"),
};

// selection: the chosen photo; quote: the plan for it (a new quote never cancels a selection still being read)
const state = { file: null, info: null, thumbUrl: null, factor: 1, plan: null, selection: 0, quote: 0, view: 0, presenting: false, examples: [], shown: [], tick: null };
const PICK = { label: els.fileLabel.textContent, meta: els.fileMeta.innerHTML };
const sizeName = (f) => (f === 1 ? "Original size" : `${f}×`);
const phone = () => matchMedia("(max-width: 720px)").matches;

// ---------------------------------------------------------------------------- the two frames share one view

class Diptych {
  constructor(root, wrap) {
    this.root = root;
    this.wrap = wrap;
    this.loading = root.querySelector(".loading");
    this.panes = [...root.querySelectorAll(".pane-frame")].map((frame) => ({ frame, img: frame.querySelector("img") }));
    this.W = 0;
    this.H = 0;
    this.s = 1;
    this.x = 0;
    this.y = 0;
    this.zoom = "fit";
    this.token = 0;
    this.pointers = new Map();
    this.onZoom = () => {};
    for (const { frame } of this.panes) this.bind(frame);
    new ResizeObserver(() => this.W && this.layout()).observe(root);
    addEventListener("resize", () => this.W && this.layout());
  }

  frame() {
    const f = this.panes[0].frame;
    return { w: f.clientWidth, h: f.clientHeight };
  }

  fitScale() {
    const { w, h } = this.frame();
    return Math.min(w / this.W, h / this.H);
  }

  /** Frames follow the photo: two frames of its aspect ratio, as large as the row and the window allow. */
  layout() {
    const gap = phone() ? 8 : 10;
    const total = this.root.clientWidth;
    const maxH = Math.min(Math.max(320, innerHeight * (phone() ? 0.62 : 0.74)), 860);
    let fw = phone() ? total : (total - gap) / 2;
    let fh = (fw * this.H) / this.W;
    if (fh > maxH) {
      fh = maxH;
      fw = (fh * this.W) / this.H;
    }
    fh = Math.max(fh, phone() ? 200 : 240);
    this.wrap.style.setProperty("--fw", `${Math.floor(fw)}px`);
    this.wrap.style.setProperty("--fh", `${Math.round(fh)}px`);
    this.wrap.style.setProperty("--gap", `${gap}px`);
    if (this.zoom === "fit") this.setZoom("fit", false);
    else this.apply();
  }

  async show({ before, after, width, height }) {
    const token = ++this.token;
    this.loading.hidden = false;
    // Each call loads into its own images; only the newest one swaps them in, together with the size.
    const load = (src, like) => {
      const img = new Image();
      img.alt = like.alt;
      img.draggable = false;
      if (!src) return Promise.resolve(img);
      // A rejected decode() is a failed display (the run stays saved for Load again), not something to show.
      return new Promise((resolve, reject) => {
        img.onload = () => resolve();
        img.onerror = () => reject(new Error("load"));
        img.src = src;
      })
        .then(() => img.decode?.())
        .then(() => img);
    };
    let pair;
    try {
      pair = await Promise.all([load(before, this.panes[0].img), load(after, this.panes[1].img)]);
    } catch {
      if (token !== this.token) return "superseded";
      this.loading.hidden = true;
      return "error";
    }
    if (token !== this.token) return "superseded";
    pair.forEach((img, i) => {
      this.panes[i].img.replaceWith(img);
      this.panes[i].img = img;
    });
    this.loading.hidden = true;
    this.panes[0].img.hidden = !before;
    els.emptyNote.hidden = !!before;
    this.W = width;
    this.H = height;
    this.zoom = "fit";
    this.layout();
    this.onZoom("fit");
    return "shown";
  }

  setZoom(z, announce = true) {
    if (!this.W) return;
    this.zoom = z;
    this.s = z === "fit" ? this.fitScale() : Number(z);
    if (z === "fit") {
      this.x = this.W / 2;
      this.y = this.H / 2;
    }
    this.apply();
    if (announce) this.onZoom(z);
  }

  zoomAt(scale, sx, sy) {
    const fit = this.fitScale();
    const next = Math.min(8, Math.max(Math.min(fit, 1), scale));
    const { w, h } = this.frame();
    const px = this.x + (sx - w / 2) / this.s;
    const py = this.y + (sy - h / 2) / this.s;
    this.s = next;
    this.x = px - (sx - w / 2) / this.s;
    this.y = py - (sy - h / 2) / this.s;
    this.zoom = Math.abs(next - 1) < 1e-6 ? "1" : Math.abs(next - 2) < 1e-6 ? "2" : "custom";
    this.apply();
    this.onZoom(this.zoom);
  }

  clamp() {
    const { w, h } = this.frame();
    const hw = w / (2 * this.s);
    const hh = h / (2 * this.s);
    this.x = this.W * this.s <= w ? this.W / 2 : Math.min(this.W - hw, Math.max(hw, this.x));
    this.y = this.H * this.s <= h ? this.H / 2 : Math.min(this.H - hh, Math.max(hh, this.y));
  }

  apply() {
    this.clamp();
    const { w, h } = this.frame();
    const tx = Math.round(w / 2 - this.x * this.s);
    const ty = Math.round(h / 2 - this.y * this.s);
    for (const { img } of this.panes) {
      img.style.width = `${this.W * this.s}px`;
      img.style.height = `${this.H * this.s}px`;
      img.style.transform = `translate(${tx}px, ${ty}px)`;
      img.style.imageRendering = this.s > 1.5 ? "pixelated" : "auto";
    }
  }

  bind(frame) {
    const usable = (e) => e.pointerType === "mouse" || els.viewerWrap.dataset.inspect === "true";
    frame.addEventListener("pointerdown", (e) => {
      if (!this.W || !usable(e)) return;
      frame.setPointerCapture(e.pointerId);
      this.pointers.set(e.pointerId, { x: e.clientX, y: e.clientY });
      frame.classList.add("is-dragging");
    });
    frame.addEventListener("pointermove", (e) => {
      const prev = this.pointers.get(e.pointerId);
      if (!prev) return;
      const now = { x: e.clientX, y: e.clientY };
      if (this.pointers.size === 1) {
        this.x -= (now.x - prev.x) / this.s;
        this.y -= (now.y - prev.y) / this.s;
        this.pointers.set(e.pointerId, now);
        if (this.zoom === "fit") {
          this.zoom = "custom";
          this.onZoom("custom");
        }
        this.apply();
        return;
      }
      const [a, b] = [...this.pointers.values()];
      const before = Math.hypot(a.x - b.x, a.y - b.y);
      this.pointers.set(e.pointerId, now);
      const [c, d] = [...this.pointers.values()];
      const after = Math.hypot(c.x - d.x, c.y - d.y);
      if (before > 0) {
        const r = frame.getBoundingClientRect();
        this.zoomAt(this.s * (after / before), (c.x + d.x) / 2 - r.left, (c.y + d.y) / 2 - r.top);
      }
    });
    const end = (e) => {
      this.pointers.delete(e.pointerId);
      if (!this.pointers.size) frame.classList.remove("is-dragging");
    };
    frame.addEventListener("pointerup", end);
    frame.addEventListener("pointercancel", end);
    // a plain wheel scrolls the page; Ctrl/⌘ + wheel (and trackpad pinch) zooms
    frame.addEventListener(
      "wheel",
      (e) => {
        if (!this.W || !(e.ctrlKey || e.metaKey)) return;
        e.preventDefault();
        const r = frame.getBoundingClientRect();
        this.zoomAt(this.s * Math.exp(-e.deltaY * 0.004), e.clientX - r.left, e.clientY - r.top);
      },
      { passive: false },
    );
    frame.addEventListener("dblclick", (e) => {
      const r = frame.getBoundingClientRect();
      if (this.zoom === "fit") this.zoomAt(1, e.clientX - r.left, e.clientY - r.top);
      else this.setZoom("fit");
    });
    // keyboard: arrows pan both frames (Shift: bigger steps), + and - zoom, Home fits
    frame.tabIndex = 0;
    frame.setAttribute("role", "group");
    frame.setAttribute("aria-label", "Arrow keys move both frames, + and - zoom, Home fits");
    frame.addEventListener("keydown", (e) => {
      if (!this.W || e.altKey || e.ctrlKey || e.metaKey) return;
      const step = (e.shiftKey ? 160 : 40) / this.s;
      const moves = { ArrowLeft: [-step, 0], ArrowRight: [step, 0], ArrowUp: [0, -step], ArrowDown: [0, step] };
      const { w, h } = this.frame();
      if (moves[e.key]) {
        this.x += moves[e.key][0];
        this.y += moves[e.key][1];
        this.apply();
      } else if (e.key === "+" || e.key === "=") this.zoomAt(this.s * 1.5, w / 2, h / 2);
      else if (e.key === "-" || e.key === "_") this.zoomAt(this.s / 1.5, w / 2, h / 2);
      else if (e.key === "Home") this.setZoom("fit");
      else return;
      e.preventDefault();
    });
  }
}

const viewer = new Diptych($("#viewer"), els.viewerWrap);
viewer.onZoom = (z) => document.querySelectorAll("[data-zoom]").forEach((b) => b.setAttribute("aria-pressed", String(b.dataset.zoom === z)));

// ---------------------------------------------------------------------------- status: one line

const demo = new Demo({
  slug: "photo-enhancer",
  onQuota: () => update(),
  turnstileEl: "#turnstile",
  turnstileTheme: () => {
    const t = document.documentElement.dataset.theme;
    if (t === "dark" || t === "light") return t;
    return matchMedia("(prefers-color-scheme: dark)").matches ? "dark" : "light";
  },
  onPhase: (phase, info) => {
    update();
    const text = { preparing: "Preparing…", uploading: "Uploading…", queued: "Waiting for a GPU…", running: "Enhancing…", loading: "Loading result…" }[phase] ?? "Working…";
    status({ busy: true, text, at: info.at ?? demo.active?.at, note: info.cold ? "The first run after a quiet spell can take 2–3 minutes." : "" });
  },
});

function status({ busy = false, error = false, text, note = "", actions = [], at = null }) {
  clearInterval(state.tick);
  els.status.hidden = false;
  els.status.dataset.state = error ? "error" : busy ? "busy" : "idle";
  const acts = actions.length
    ? `<span class="actions">${actions
        .map((a, i) =>
          a.href
            ? `<a href="${a.href}" target="_blank" rel="noopener">${a.label}</a>`
            : `<button type="button" class="${a.primary ? "primary" : ""}" data-i="${i}">${a.label}</button>`,
        )
        .join("")}</span>`
    : "";
  els.status.innerHTML = `<span class="status-text">${busy ? '<span class="spinner" aria-hidden="true"></span>' : ""}<strong>${text}</strong>${note ? ` <span class="meta">${note}</span>` : ""}<span class="meta elapsed"></span></span>${acts}`;
  els.status.querySelectorAll("button[data-i]").forEach((b) => b.addEventListener("click", () => actions[Number(b.dataset.i)].run()));
  if (at) {
    const el = els.status.querySelector(".elapsed");
    const upd = () => (el.textContent = ` · ${Math.round((Date.now() - at) / 1000)}${NBSP}s`);
    upd();
    state.tick = setInterval(upd, 1000);
  }
}

function showError(err) {
  const [title, note] = errorCopy(err, demo, COPY);
  if (INFO_KINDS.has(err.kind) && !demo.active?.op) return status({ text: title, note });
  const actions = [];
  if (isResumable(err) && demo.active) actions.push({ label: "Resume", primary: true, run: () => go(() => demo.continue()) });
  if (err.kind === "download" && demo.active?.op) actions.push({ label: "Load again", primary: true, run: () => go(() => demo.continue()) });
  if (demo.active && actions.length)
    actions.push({
      label: "Start over",
      run: () => {
        demo.forget();
        els.status.hidden = true;
        update();
      },
    });
  status({ error: true, text: title, note, actions });
}

// ---------------------------------------------------------------------------- photo and plan

async function selectFile(file) {
  if (!file || demo.busy) return;
  const sel = ++state.selection;
  state.quote++; // a quote still running is for the previous photo
  state.file = null;
  state.info = null;
  state.plan = null;
  demo.dropPrepared();
  update();
  if (!demo.active) els.status.hidden = true;
  let info;
  try {
    info = await demo.inspect(file);
  } catch (err) {
    if (sel !== state.selection) return;
    clearPhoto();
    return planError(err);
  }
  const thumb = await thumbnail(file, info, 120);
  if (sel !== state.selection) return thumb && URL.revokeObjectURL(thumb);
  if (state.thumbUrl) URL.revokeObjectURL(state.thumbUrl);
  state.file = file;
  state.info = info;
  state.thumbUrl = thumb;
  els.thumb.hidden = !thumb;
  if (thumb) els.thumb.src = thumb;
  els.fileLabel.textContent = "Change photo";
  els.fileMeta.innerHTML = `<strong>${escapeHtml(file.name || "Pasted photo")}</strong> · <span class="numeric">${dims(info.width, info.height)}</span>`;
  refreshPlan();
}

/** No photo: the file, its size, its plan and its preview go together. */
function clearPhoto() {
  if (state.thumbUrl) URL.revokeObjectURL(state.thumbUrl);
  Object.assign(state, { file: null, info: null, plan: null, thumbUrl: null });
  els.thumb.hidden = true;
  els.thumb.removeAttribute("src");
  els.fileLabel.textContent = PICK.label;
  els.fileMeta.innerHTML = PICK.meta;
}

function planError(err, current = () => true) {
  if (!current()) return;
  const [title, note] = errorCopy(err, demo, COPY);
  els.plan.innerHTML = `<span class="plan-error"><strong>${title}</strong></span> ${note}`;
  state.plan = null;
  update();
}

function applyPlan(plan) {
  state.plan = plan;
  const price = demo.config.price;
  const api = plan.chips ? ` · via the API ${plan.chips}${NBSP}chip${plan.chips === 1 ? "" : "s"}${price ? ` (${usd(plan.chips * price.min)}–${usd(plan.chips * price.max)})` : ""}` : "";
  const reduced = plan.resized ? ` · reduced from ${dims(plan.source.width, plan.source.height)}` : "";
  els.plan.textContent = `Output ${dims(plan.out_w, plan.out_h)}${reduced}${api}`;
}

async function refreshPlan() {
  const { file, info } = state;
  if (!file || !info) return;
  if (demo.loginNeeded()) demo.noteLoginWall();
  const sel = state.selection;
  const q = ++state.quote;
  const current = () => q === state.quote && sel === state.selection;
  state.plan = null;
  update();
  if (!demo.config) return; // start() quotes again once the service answers
  try {
    const plan = await demo.quote(info.width, info.height, state.factor);
    if (!current()) return;
    applyPlan(plan);
  } catch (err) {
    planError(err, current);
  }
  update();
}

function update() {
  const login = demo.loginNeeded();
  const limited = demo.limited();
  const working = demo.busy || state.presenting;
  const pending = !!demo.active && !working;
  els.run.disabled = working || (!pending && (!state.plan || demo.configFailed || limited));
  els.run.textContent = working ? "Enhancing…" : pending ? "Resume" : login ? "Log in to enhance" : limited ? "No runs left today" : "Enhance";
}

// ---------------------------------------------------------------------------- run

async function go(step) {
  update();
  try {
    const done = await step();
    state.presenting = true;
    update();
    await present(done);
  } catch (err) {
    const e = err instanceof DemoError ? err : new DemoError("server");
    // The browser could prepare only a smaller image: show that plan; the next press sends it.
    if (e.kind === "reduced" && e.plan && e.file === state.file && e.from === state.plan) applyPlan(e.plan);
    showError(e);
  } finally {
    state.presenting = false;
    update();
  }
}

function setBeforeLabel(factor) {
  const text = factor === 1 ? "Before" : `Before · resized ${factor}×`;
  els.beforeLabel.textContent = text;
  viewer.panes[0].frame.closest(".pane").dataset.label = text;
}

async function present(done) {
  const view = ++state.view;
  const op = done.op;
  setBeforeLabel(op.factor);
  const outcome = await viewer.show({ before: done.beforeUrl, after: done.afterUrl, width: op.out_w, height: op.out_h });
  if (outcome !== "shown" || view !== state.view) {
    // Not on screen: the run stays saved, so it can be loaded again (Load again, Resume, or after a reload).
    URL.revokeObjectURL(done.afterUrl);
    if (outcome === "error") throw new DemoError("download");
    if (demo.active?.op?.op === op.op)
      status({ text: "Your result is ready.", actions: [{ label: "Show it", primary: true, run: () => go(() => demo.continue()) }] });
    return;
  }
  demo.shown(done);
  for (const u of state.shown) URL.revokeObjectURL(u);
  state.shown = [done.afterUrl, done.beforeUrl].filter(Boolean);
  markExample(null);
  els.caption.innerHTML = `<strong>Your photo</strong><span class="meta numeric">${sizeName(op.factor)} · ${dims(op.out_w, op.out_h)}</span>`;
  const left = demo.runsLeft();
  const actions = [{ label: "Open megapotato.app ↗", href: demo.site("/tool", "get-png") }];
  if (phone())
    actions.push({ label: "View result", run: () => els.viewerWrap.scrollIntoView({ behavior: matchMedia("(prefers-reduced-motion: reduce)").matches ? "auto" : "smooth" }) });
  status({
    text: `Ready${left !== null ? ` · ${left} run${left === 1 ? "" : "s"} left today` : ""}`,
    note: "Upload your image on megapotato.app to download it.",
    actions,
  });
}

// ---------------------------------------------------------------------------- examples

async function loadExamples() {
  try {
    state.examples = await (await fetch("examples/examples.json")).json();
  } catch {
    state.examples = [];
  }
  els.examples.innerHTML = state.examples
    .map(
      (ex, i) => `<li style="--ar:${(ex.in[0] / ex.in[1]).toFixed(4)}"><button type="button" class="example" data-i="${i}" aria-pressed="false">
        <span class="example-image"><img src="${escapeHtml(ex.thumb)}" alt="" loading="lazy"></span>
        <span class="example-text"><span>${escapeHtml(ex.title)}</span><small>${sizeName(ex.factor)}</small></span>
      </button></li>`,
    )
    .join("");
  els.credits.innerHTML = state.examples
    .map((ex) => `<li>${escapeHtml(ex.credit)} <a href="${escapeHtml(ex.source)}" target="_blank" rel="noopener">Source</a>. ${escapeHtml(ex.licence ?? "")}</li>`)
    .join("");
  els.examples.querySelectorAll(".example").forEach((b) => b.addEventListener("click", () => showExample(Number(b.dataset.i))));
}

function markExample(index) {
  els.examples.querySelectorAll(".example").forEach((b, i) => b.setAttribute("aria-pressed", String(i === index)));
}

async function showExample(i) {
  const ex = state.examples[i];
  if (!ex) return;
  const view = ++state.view;
  markExample(i);
  setBeforeLabel(ex.factor);
  els.caption.innerHTML = `<strong>${escapeHtml(ex.title)}</strong><span class="meta numeric">${escapeHtml(ex.caption)} · ${sizeName(ex.factor)} · ${dims(ex.out[0], ex.out[1])}</span>`;
  const outcome = await viewer.show({ before: ex.input, after: ex.result, width: ex.out[0], height: ex.out[1] });
  if (view === state.view && outcome === "shown") {
    for (const u of state.shown) URL.revokeObjectURL(u);
    state.shown = [];
  }
}

// ---------------------------------------------------------------------------- wiring

function wire() {
  els.file.addEventListener("change", () => selectFile(els.file.files?.[0]));
  // the whole page is the drop target
  let depth = 0;
  const hasFiles = (e) => [...(e.dataTransfer?.types ?? [])].includes("Files");
  document.addEventListener("dragenter", (e) => {
    if (!hasFiles(e)) return;
    e.preventDefault();
    depth += 1;
    document.body.dataset.dragover = "true";
  });
  document.addEventListener("dragover", (e) => {
    if (!hasFiles(e)) return;
    e.preventDefault();
  });
  document.addEventListener("dragleave", (e) => {
    if (!hasFiles(e)) return;
    depth = Math.max(0, depth - 1);
    if (!depth) delete document.body.dataset.dragover;
  });
  document.addEventListener("drop", (e) => {
    if (!hasFiles(e)) return;
    e.preventDefault();
    depth = 0;
    delete document.body.dataset.dragover;
    selectFile(e.dataTransfer?.files?.[0]);
  });
  document.addEventListener("paste", (e) => {
    const item = [...(e.clipboardData?.items ?? [])].find((it) => it.kind === "file" && it.type.startsWith("image/"));
    if (item) selectFile(item.getAsFile());
  });
  document.querySelectorAll('input[name="mode"]').forEach((r) =>
    r.addEventListener("change", () => {
      state.factor = Number(r.value);
      refreshPlan();
    }),
  );
  document.querySelectorAll("[data-zoom]").forEach((b) => b.addEventListener("click", () => viewer.setZoom(b.dataset.zoom)));
  els.inspect.addEventListener("click", () => {
    const on = els.viewerWrap.dataset.inspect !== "true";
    els.viewerWrap.dataset.inspect = String(on);
    els.inspect.setAttribute("aria-pressed", String(on));
    els.inspect.textContent = on ? "Done" : "Inspect";
  });
  els.run.addEventListener("click", () => {
    if (demo.active) return go(() => demo.continue());
    if (state.file && state.plan) go(() => demo.start(state.file, state.plan));
  });
  viewer.panes[1].frame.closest(".pane").dataset.label = "After";
}

async function start() {
  wire();
  const saved = demo.hasSaved();
  await loadExamples();
  if (state.examples.length) showExample(0);
  const c = await demo.init();
  if (!c) {
    els.notice.hidden = false;
    els.notice.textContent = "Live enhancing is unavailable right now. The examples still work.";
    return update();
  }
  els.quota.textContent = c.login_required
    ? `Preview only · ${c.jobs_per_day} free runs a day · Log in to Hugging Face (top right) to run`
    : `${c.jobs_per_day} free runs a day`;
  update();
  refreshPlan(); // a photo chosen while the page was still starting gets its plan now
  if (!saved && demo.loginNeeded()) showError(new DemoError("login"));
  if (saved) {
    status({ busy: true, text: "Picking up your last run…" });
    await go(() => demo.resume().then((r) => r ?? Promise.reject(new DemoError("unconfirmed"))));
  }
}

start();