import { Client } from "https://cdn.jsdelivr.net/npm/@gradio/client@2.7.1/dist/index.min.js";

const BACKEND = "SupraLabs/supra2-img-demo";
const BACKEND_URL = "https://supralabs-supra2-img-demo.hf.space";
const STEPS = { Fast: 3, Quality: 8, "Ultra Design": 16 };

const form = document.querySelector("#generator-form");
const prompt = document.querySelector("#prompt");
const generateButton = document.querySelector("#generate");
const buttonLabel = document.querySelector("#button-label");
const guidance = document.querySelector("#guidance");
const guidanceValue = document.querySelector("#guidance-value");
const seed = document.querySelector("#seed");
const randomize = document.querySelector("#randomize");
const image = document.querySelector("#result-image");
const placeholder = document.querySelector("#placeholder");
const download = document.querySelector("#download");
const seedUsed = document.querySelector("#seed-used");
const resultState = document.querySelector("#result-state");
const status = document.querySelector("#status");
const modeLabel = document.querySelector("#selected-mode");
const modeAvailability = document.querySelector("#mode-availability");
const level = document.querySelector("#level");
const levelTrack = document.querySelector("#level-track");
const levelNodes = [...document.querySelectorAll(".level-node")];
const levelLabels = [...document.querySelectorAll(".level-labels button")];
const progressPanel = document.querySelector("#render-progress");
const progressTrack = document.querySelector("#progress-track");
const progressFill = document.querySelector("#progress-fill");
const progressPercent = document.querySelector("#progress-percent");
const progressStage = document.querySelector("#progress-stage");
const LEVELS = ["Fast", "Quality", "Ultra Design"];

let mode = "Quality";
let clientPromise;
let dragging = false;
let levelTouched = false;
let shownProgress = 0;
let displayedProgress = 0;
let progressFrame = 0;
let backendMinimumSteps = 10;

function setStatus(message, kind = "") {
  status.textContent = message;
  status.className = `status ${kind}`.trim();
}

function setMode(nextMode) {
  if (!(nextMode in STEPS)) return;
  const previous = mode;
  mode = nextMode;
  level.dataset.level = mode;
  modeLabel.textContent = `${mode} · ${STEPS[mode]} steps`;
  levelTrack.setAttribute("aria-valuenow", String(LEVELS.indexOf(mode)));
  levelTrack.setAttribute("aria-valuetext", `${mode}, ${STEPS[mode]} steps`);
  modeAvailability.textContent = STEPS[mode] < backendMinimumSteps
    ? `${mode} can be selected, but generation at ${STEPS[mode]} steps is awaiting the CPU service update.`
    : "";
  if (mode === "Ultra Design" && previous !== mode) requestAnimationFrame(playUltra);
  if (mode !== "Ultra Design" && previous === "Ultra Design") cancelUltra();
}

function cancelUltra() {
  level.querySelectorAll(".level-rainbow, .level-node").forEach((item) =>
    item.getAnimations().forEach((animation) => { if (animation.id === "rainbow-grow") animation.cancel(); }));
}

function playUltra() {
  cancelUltra();
  if (matchMedia("(prefers-reduced-motion: reduce)").matches) return;
  const rainbow = level.querySelector(".level-rainbow");
  const knob = level.querySelector('.level-node[data-mode="Ultra Design"]');
  const grow = rainbow.animate(
    [{ width: "0%", opacity: 1 }, { width: "100%", opacity: 1 }],
    { duration: 1500, delay: 520, easing: "cubic-bezier(.45,.05,.25,1)", fill: "backwards" }
  );
  grow.id = "rainbow-grow";
  const pulse = knob.animate(
    [{ transform: "scale(1)" }, { transform: "scale(1.9)", offset: .55 }, { transform: "scale(1.65)" }],
    { duration: 520, delay: 1900, easing: "ease-out", fill: "backwards" }
  );
  pulse.id = "rainbow-grow";
}

function setProgress(percent, stage) {
  shownProgress = Math.max(shownProgress, Math.min(100, Math.max(0, Math.round(percent))));
  progressPanel.hidden = false;
  progressFill.style.width = `${shownProgress}%`;
  progressStage.textContent = stage;
  progressTrack.setAttribute("aria-valuenow", String(shownProgress));
  cancelAnimationFrame(progressFrame);
  if (matchMedia("(prefers-reduced-motion: reduce)").matches) {
    displayedProgress = shownProgress;
    progressPercent.textContent = `${shownProgress}%`;
    return;
  }
  const start = displayedProgress;
  const target = shownProgress;
  const started = performance.now();
  const tick = (now) => {
    const t = Math.min(1, (now - started) / 700);
    displayedProgress = start + (target - start) * (1 - (1 - t) ** 3);
    progressPercent.textContent = `${Math.round(displayedProgress)}%`;
    if (t < 1) progressFrame = requestAnimationFrame(tick);
  };
  progressFrame = requestAnimationFrame(tick);
}

function pickLevel(clientX) {
  const bounds = levelTrack.getBoundingClientRect();
  const fraction = Math.min(1, Math.max(0, (clientX - bounds.left - 18) / (bounds.width - 36)));
  setMode(LEVELS[Math.round(fraction * 2)]);
}

async function refreshAvailableModes() {
  try {
    const response = await fetch(`${BACKEND_URL}/config`, { cache: "no-store" });
    if (!response.ok) throw new Error("CPU backend unavailable");
    const config = await response.json();
    const stepControl = config.components?.find((part) => part.type === "slider" && part.props?.label === "Sampling steps");
    backendMinimumSteps = Number(stepControl?.props?.minimum ?? 10);
    if (!levelTouched && STEPS[mode] < backendMinimumSteps) setMode("Ultra Design");
    else setMode(mode);
  } catch {
    modeAvailability.textContent = "The CPU service is waking up. You can try generating an image shortly.";
  }
}

function setBusy(busy) {
  generateButton.disabled = busy;
  buttonLabel.textContent = busy ? "Generating..." : "Generate image";
  resultState.textContent = busy ? "Working on CPU" : "Ready to create";
}

function imageUrl(value) {
  if (typeof value === "string") return value;
  if (value && typeof value.url === "string") return value.url;
  if (value && typeof value.path === "string") {
    return new URL(`/gradio_api/file=${value.path}`, "https://supralabs-supra2-img-demo.hf.space").href;
  }
  throw new Error("The image service returned no image URL.");
}

for (const option of levelLabels) option.addEventListener("click", () => { levelTouched = true; setMode(option.dataset.mode); });
levelTrack.addEventListener("pointerdown", (event) => {
  dragging = true;
  levelTouched = true;
  levelTrack.focus({ preventScroll: true });
  try { levelTrack.setPointerCapture(event.pointerId); } catch { /* Pointer capture is optional. */ }
  pickLevel(event.clientX);
  event.preventDefault();
});
levelTrack.addEventListener("pointermove", (event) => { if (dragging) pickLevel(event.clientX); });
levelTrack.addEventListener("pointerup", () => { dragging = false; });
levelTrack.addEventListener("pointercancel", () => { dragging = false; });
levelTrack.addEventListener("keydown", (event) => {
  levelTouched = true;
  const index = LEVELS.indexOf(mode);
  const next = event.key === "ArrowRight" || event.key === "ArrowUp" ? index + 1
    : event.key === "ArrowLeft" || event.key === "ArrowDown" ? index - 1
    : event.key === "Home" ? 0 : event.key === "End" ? 2 : null;
  if (next !== null) { event.preventDefault(); setMode(LEVELS[Math.max(0, Math.min(2, next))]); }
});
refreshAvailableModes();
guidance.addEventListener("input", () => { guidanceValue.value = Number(guidance.value).toFixed(1); });
document.querySelectorAll("#examples button").forEach((example) => {
  example.addEventListener("click", () => { prompt.value = example.textContent.trim(); prompt.focus(); });
});

form.addEventListener("submit", async (event) => {
  event.preventDefault();
  const description = prompt.value.trim();
  if (!description) { prompt.focus(); return; }
  if (STEPS[mode] < backendMinimumSteps) {
    setStatus(`${mode} needs the CPU service update before it can generate at ${STEPS[mode]} steps.`, "error");
    return;
  }
  const seedNumber = Number(seed.value);
  if (!randomize.checked && (!Number.isInteger(seedNumber) || seedNumber < 0 || seedNumber > 2147483647)) {
    setStatus("Choose a seed between 0 and 2147483647.", "error");
    seed.focus();
    return;
  }

  setBusy(true);
  setStatus("Connecting to the CPU image service...");
  shownProgress = 0;
  displayedProgress = 0;
  cancelAnimationFrame(progressFrame);
  progressPercent.textContent = "0%";
  setProgress(0, "Connecting to CPU");
  try {
    clientPromise ??= Client.connect(BACKEND, { events: ["status", "data"] });
    const client = await clientPromise;
    const job = client.submit("/generate", [description, STEPS[mode], Number(guidance.value), seedNumber, randomize.checked]);
    let gotImage = false;
    for await (const message of job) {
      if (message.type === "status") {
        if (message.stage === "pending") {
          const position = Number.isInteger(message.position) ? ` · position ${message.position + 1}` : "";
          setStatus(`Queued for CPU generation${position}`);
          setProgress(0, "Waiting in queue");
        } else if (message.stage === "generating") {
          const progress = message.progress_data?.find((part) => part?.desc);
          setStatus(progress?.desc || "Generating your image on CPU...");
          let fraction = Number(progress?.progress);
          if (!Number.isFinite(fraction) && Number(progress?.length) > 0) fraction = Number(progress.index) / Number(progress.length);
          if (!Number.isFinite(fraction) && progress?.desc) {
            const steps = progress.desc.match(/(\d+)\s*\/\s*(\d+)/);
            if (steps) fraction = Number(steps[1]) / Number(steps[2]);
          }
          if (Number.isFinite(fraction)) setProgress(4 + 90 * Math.min(1, Math.max(0, fraction)), progress?.desc || "Rendering image");
          else setProgress(4, "Preparing image");
        } else if (message.stage === "error") {
          throw new Error(message.message || "Image generation failed.");
        }
      }
      if (message.type === "data" && Array.isArray(message.data)) {
        const url = imageUrl(message.data[0]);
        image.src = url;
        image.hidden = false;
        placeholder.hidden = true;
        download.href = url;
        download.hidden = false;
        seedUsed.textContent = `Seed ${message.data[1]}`;
        setProgress(100, "Image ready");
        gotImage = true;
      }
    }
    if (!gotImage) throw new Error("The image service did not return an image.");
    setStatus(`${mode} complete · ${STEPS[mode]} steps`, "success");
  } catch (error) {
    console.error(error);
    clientPromise = undefined;
    setStatus(error instanceof Error ? error.message : "Image generation failed. Please try again.", "error");
  } finally {
    setBusy(false);
  }
});