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

const $ = (selector, root = document) => root.querySelector(selector);
const $$ = (selector, root = document) => [...root.querySelectorAll(selector)];
const dims = [64, 128, 256, 512, 1024, 2048, 4096];
const colors = ["#d9ff63", "#66e3cf", "#ff9c62", "#a993ff"];

const state = {
  client: null,
  connecting: null,
  searchImage: null,
  searchVideo: null,
  candidateMedia: [],
  pairQueryImage: null,
  pairQueryVideo: null,
  pairCandidateImage: null,
  pairCandidateVideo: null,
  searchDimension: 1024,
  compareDimension: 1024,
  objectUrls: new Map(),
  jobTimer: null,
  jobStarted: 0,
};

const examples = {
  llama: {
    text: "Which Llama 4 model variants are available?",
    dimension: 512,
  },
  tofu: {
    text: "How is mapo tofu prepared?",
    dimension: 1024,
  },
  safety: {
    text: "Find the environmental assessment page about driver training and temporary road closures.",
    dimension: 256,
  },
  image: {
    text: "Match this screenshot to the most relevant description.",
    asset: "/assets/llama4_hgf.png",
    filename: "llama4_hgf.png",
    type: "image/png",
    custom:
      "Llama family :: Scout and Maverick are multimodal mixture-of-experts model variants.\nRecipe :: Soft tofu simmered in spicy chili-bean sauce.",
    include: false,
    dimension: 256,
  },
  video: {
    text: "What dish is being prepared in this clip?",
    asset: "/assets/mapo_tofu.mp4",
    filename: "mapo_tofu.mp4",
    type: "video/mp4",
    custom:
      "Sichuan classic :: Mapo tofu combines soft tofu with a spicy, numbing bean-paste sauce.\nSpaceflight :: A launch vehicle carries a satellite into orbit.",
    include: false,
    dimension: 512,
  },
};

function escapeHtml(value) {
  return String(value ?? "")
    .replaceAll("&", "&")
    .replaceAll("<", "&lt;")
    .replaceAll(">", "&gt;")
    .replaceAll('"', "&quot;")
    .replaceAll("'", "&#039;");
}

function formatScore(value) {
  const number = Number(value);
  return `${number >= 0 ? "+" : ""}${number.toFixed(3)}`;
}

function formatSeconds(value) {
  return `${Number(value).toFixed(1)}s`;
}

function clearObjectUrls() {
  for (const value of state.objectUrls.values()) URL.revokeObjectURL(value);
  state.objectUrls.clear();
}

async function connectClient() {
  if (state.client) return state.client;
  if (state.connecting) return state.connecting;
  const status = $("#model-status");
  state.connecting = (async () => {
    try {
      const client = await Client.connect(window.location.origin);
      await client.view_api();
      state.client = client;
      status.className = "model-status ready";
      $("span", status).textContent = "Model ready";
      return client;
    } catch (error) {
      status.className = "model-status error";
      $("span", status).textContent = "Connection failed";
      state.connecting = null;
      throw error;
    }
  })();
  return state.connecting;
}

function switchMode(mode) {
  $$(".nav-pill").forEach((button) => button.classList.toggle("active", button.dataset.mode === mode));
  $$("[data-panel]").forEach((panel) => { panel.hidden = panel.dataset.panel !== mode; });
  $(mode === "search" ? '[data-panel="search"]' : '[data-panel="compare"]').scrollIntoView({ behavior: "smooth", block: "start" });
}

function setDimension(group, value) {
  const dimension = Number(value);
  state[`${group}Dimension`] = dimension;
  $(`[data-dimension-group="${group}"]`).querySelectorAll("button").forEach((button) => {
    button.classList.toggle("active", Number(button.dataset.value) === dimension);
  });
  $(`#${group}-dimension`).value = String(dimension);
  const output = $(`#${group}-dimension-output`);
  if (output) output.textContent = `${dimension.toLocaleString()}D`;
}

function fileLabel(drop, file) {
  const chip = $(".file-chip", drop);
  if (!chip) return;
  chip.hidden = !file;
  if (file) $("span", chip).textContent = file.name;
}

function configureDropzone(dropSelector, inputSelector, stateKey, otherStateKey, otherDropSelector) {
  const drop = $(dropSelector);
  const input = $(inputSelector);
  const acceptFile = (file) => {
    if (!file) return;
    state[stateKey] = file;
    fileLabel(drop, file);
    if (otherStateKey) {
      state[otherStateKey] = null;
      fileLabel($(otherDropSelector), null);
      const otherInput = $(`${otherDropSelector} input`);
      if (otherInput) otherInput.value = "";
    }
  };
  input.addEventListener("change", () => acceptFile(input.files[0]));
  ["dragenter", "dragover"].forEach((eventName) => drop.addEventListener(eventName, (event) => {
    event.preventDefault();
    drop.classList.add("dragover");
  }));
  ["dragleave", "drop"].forEach((eventName) => drop.addEventListener(eventName, (event) => {
    event.preventDefault();
    drop.classList.remove("dragover");
  }));
  drop.addEventListener("drop", (event) => acceptFile(event.dataTransfer.files[0]));
  const remove = $(".file-chip button", drop);
  if (remove) remove.addEventListener("click", (event) => {
    event.preventDefault();
    event.stopPropagation();
    input.value = "";
    state[stateKey] = null;
    fileLabel(drop, null);
  });
}

function renderCandidateChips() {
  const root = $("#candidate-chips");
  root.innerHTML = state.candidateMedia.map((file, index) => `
    <span class="candidate-chip"><span>${escapeHtml(file.name)}</span><button type="button" data-remove-media="${index}" aria-label="Remove ${escapeHtml(file.name)}">×</button></span>
  `).join("");
  $$('[data-remove-media]', root).forEach((button) => button.addEventListener("click", () => {
    state.candidateMedia.splice(Number(button.dataset.removeMedia), 1);
    renderCandidateChips();
  }));
}

function bindPairFile(inputSelector, stateKey, otherKey, labelSelector) {
  const input = $(inputSelector);
  input.addEventListener("change", () => {
    const file = input.files[0] ?? null;
    state[stateKey] = file;
    if (file && otherKey) {
      state[otherKey] = null;
      const side = input.closest(".compare-side");
      const otherInput = input.accept.includes("image") ? $('input[accept="video/*"]', side) : $('input[accept="image/*"]', side);
      if (otherInput) otherInput.value = "";
    }
    const label = $(labelSelector);
    label.hidden = !file;
    label.textContent = file ? file.name : "";
  });
}

function startJob(mode, heavy = false) {
  const toast = $("#job-toast");
  const phase = $("#job-phase");
  const elapsed = $("#job-elapsed");
  const progress = $("#job-progress-bar");
  $("#job-label").textContent = mode === "search" ? "RESOLVING SEMANTIC FIELD" : "MEASURING VECTOR ALIGNMENT";
  phase.textContent = "Waiting for the ZeroGPU allocator…";
  progress.style.width = "3%";
  toast.hidden = false;
  state.jobStarted = Date.now();
  clearInterval(state.jobTimer);
  state.jobTimer = setInterval(() => {
    const seconds = Math.floor((Date.now() - state.jobStarted) / 1000);
    elapsed.textContent = `${String(Math.floor(seconds / 60)).padStart(2, "0")}:${String(seconds % 60).padStart(2, "0")} elapsed`;
    const first = heavy ? 13 : 7;
    if (seconds < 3) {
      phase.textContent = "Waiting for the ZeroGPU allocator…";
      progress.style.width = "8%";
    } else if (seconds < first) {
      phase.textContent = "Encoding inputs into 4,096 dimensions…";
      progress.style.width = `${Math.min(52, 18 + seconds * 4)}%`;
    } else {
      phase.textContent = mode === "search" ? "Ranking the multimodal universe…" : "Tracing the Matryoshka curve…";
      progress.style.width = `${Math.min(92, 55 + seconds)}%`;
    }
  }, 500);
}

function finishJob() {
  clearInterval(state.jobTimer);
  $("#job-phase").textContent = "Semantic field resolved.";
  $("#job-progress-bar").style.width = "100%";
  setTimeout(() => { $("#job-toast").hidden = true; }, 900);
}

function failJob(error) {
  clearInterval(state.jobTimer);
  $("#job-toast").hidden = true;
  const toast = $("#error-toast");
  $("span", toast).textContent = error?.message || String(error) || "The request could not be completed.";
  toast.hidden = false;
}

function unwrapResult(result) {
  let data = result?.data ?? result;
  if (Array.isArray(data) && data.length === 1) data = data[0];
  if (typeof data === "string") {
    try { return JSON.parse(data); } catch { return data; }
  }
  return data;
}

function mediaUrl(item) {
  if (item.media_url) return item.media_url;
  return state.objectUrls.get(item.key) ?? null;
}

function mediaMarkup(item, className = "rank-media") {
  const url = mediaUrl(item);
  if (!url) return `<span class="${className} text-tile">Aa</span>`;
  if (item.kind === "video") return `<video class="${className}" src="${escapeHtml(url)}" muted playsinline preload="metadata"></video>`;
  return `<img class="${className}" src="${escapeHtml(url)}" alt="" loading="lazy" />`;
}

function rankingMarkup(items) {
  return `<div class="ranking-list">${items.map((item) => {
    const fill = Math.max(2, Math.min(100, Math.max(0, Number(item.score)) * 100));
    return `<article class="ranking-card" style="--rank:${item.rank}">
      <span class="rank-index">${String(item.rank).padStart(2, "0")}</span>
      ${mediaMarkup(item)}
      <div class="rank-copy"><div class="rank-title"><b>${escapeHtml(item.title)}</b><span class="kind-tag">${escapeHtml(item.kind)}</span></div><p>${escapeHtml(item.description)}</p><div class="score-line"><i style="width:${fill}%"></i></div></div>
      <div class="rank-value">${formatScore(item.score)}<small>COSINE</small></div>
    </article>`;
  }).join("")}</div>`;
}

function visualMarkup(items) {
  const visual = items.filter((item) => mediaUrl(item));
  if (!visual.length) return `<div class="compare-empty"><p>No visual candidates in this result set.</p></div>`;
  return `<div class="visual-grid">${visual.map((item) => `<figure class="visual-card">
    ${item.kind === "video" ? `<video src="${escapeHtml(mediaUrl(item))}" controls muted playsinline preload="metadata"></video>` : `<img src="${escapeHtml(mediaUrl(item))}" alt="${escapeHtml(item.title)}" loading="lazy" />`}
    <figcaption><span>#${item.rank} · ${escapeHtml(item.title)}</span><span>${formatScore(item.score)}</span></figcaption>
  </figure>`).join("")}</div>`;
}

function lineChart(dimensions, series, title) {
  const width = 720;
  const height = 260;
  const pad = { l: 46, r: 15, t: 20, b: 36 };
  const values = series.flatMap((item) => item.values);
  let low = Math.max(-1, Math.min(...values) - 0.08);
  let high = Math.min(1, Math.max(...values) + 0.08);
  if (high - low < 0.2) { const mid = (high + low) / 2; low = Math.max(-1, mid - 0.1); high = Math.min(1, mid + 0.1); }
  const x = (index) => pad.l + (index * (width - pad.l - pad.r)) / (dimensions.length - 1);
  const y = (value) => pad.t + ((high - value) * (height - pad.t - pad.b)) / Math.max(0.0001, high - low);
  const grid = Array.from({ length: 5 }, (_, index) => {
    const value = high - (index * (high - low)) / 4;
    return `<line class="chart-grid-line" x1="${pad.l}" y1="${y(value)}" x2="${width-pad.r}" y2="${y(value)}"/><text class="chart-label" x="${pad.l-8}" y="${y(value)+3}" text-anchor="end">${formatScore(value).slice(0,-1)}</text>`;
  }).join("");
  const labels = dimensions.map((value, index) => `<text class="chart-label" x="${x(index)}" y="${height-12}" text-anchor="middle">${value}</text>`).join("");
  const paths = series.map((item, seriesIndex) => {
    const color = colors[seriesIndex % colors.length];
    const points = item.values.map((value, index) => `${x(index)},${y(value)}`).join(" ");
    const dots = item.values.map((value, index) => `<circle class="chart-dot" cx="${x(index)}" cy="${y(value)}" r="3" fill="${color}"/>`).join("");
    return `<polyline class="chart-line" points="${points}" stroke="${color}"/>${dots}`;
  }).join("");
  const legend = series.map((item, index) => `<span><i style="background:${colors[index % colors.length]}"></i>${escapeHtml(item.title.slice(0, 40))}</span>`).join("");
  return `<section class="viz-panel"><div class="viz-head"><div><span>MATRYOSHKA SCOPE</span><b>${escapeHtml(title)}</b></div><small>64 → 4096D</small></div><svg class="line-chart" viewBox="0 0 ${width} ${height}" role="img" aria-label="Cosine similarity across embedding dimensions">${grid}${labels}${paths}</svg><div class="chart-legend">${legend}</div></section>`;
}

function fingerprintChart(values, title, dimension) {
  const width = 720;
  const height = 155;
  const center = 78;
  const scale = Math.max(...values.map((value) => Math.abs(value)), 0.000001);
  const barWidth = (width - 16) / values.length;
  const bars = values.map((value, index) => {
    const magnitude = Math.min(64, (Math.abs(value) / scale) * 64);
    const y = value >= 0 ? center - magnitude : center;
    return `<rect x="${8 + index * barWidth}" y="${y}" width="${Math.max(1.2,barWidth-1.4)}" height="${magnitude}" rx="1" fill="${value >= 0 ? "#66e3cf" : "#ff9c62"}" opacity=".86"/>`;
  }).join("");
  return `<section class="viz-panel"><div class="viz-head"><div><span>VECTOR FINGERPRINT</span><b>${escapeHtml(title)}</b></div><small>${Number(dimension).toLocaleString()} values</small></div><svg class="fingerprint-svg" viewBox="0 0 ${width} ${height}" role="img" aria-label="Compressed signed vector fingerprint"><line class="fingerprint-zero" x1="8" y1="${center}" x2="${width-8}" y2="${center}"/>${bars}</svg><div class="fingerprint-key"><span><i style="background:#66e3cf"></i>POSITIVE</span><span><i style="background:#ff9c62"></i>NEGATIVE</span><span>96 POOLED SLICES</span></div></section>`;
}

function telemetryMarkup(data) {
  const compact = { ...data };
  delete compact.rankings;
  delete compact.dimension_series;
  delete compact.query_fingerprint;
  delete compact.candidate_fingerprint;
  return `<details class="telemetry"><summary><span>EMBEDDING TELEMETRY</span><span>INSPECT JSON ＋</span></summary><pre>${escapeHtml(JSON.stringify(compact, null, 2))}</pre></details>`;
}

function analysisMarkup(data) {
  return `<div class="viz-stack">${lineChart(data.dimensions, data.dimension_series, "Does the ranking survive compression?")}${fingerprintChart(data.query_fingerprint, `Query · ${data.query_kind}`, data.dimension)}${telemetryMarkup(data)}</div>`;
}

function bindResultTabs(root) {
  $$('[data-result-tab]', root).forEach((button) => button.addEventListener("click", () => {
    $$('[data-result-tab]', root).forEach((item) => item.classList.toggle("active", item === button));
    $$('[data-result-view]', root).forEach((view) => { view.hidden = view.dataset.resultView !== button.dataset.resultTab; });
  }));
}

function renderSearchResults(data) {
  const root = $("#search-results");
  root.innerHTML = `
    <div class="result-topbar"><div><span class="signal"><i></i>FIELD RESOLVED</span><small>tencent / WeMM-Embedding-9B</small></div><span class="result-count">${data.candidate_count} CANDIDATES · ${Number(data.dimension).toLocaleString()}D</span></div>
    <section class="summary-card"><div><span>TOP SEMANTIC MATCH</span><h3>${escapeHtml(data.top_match.title)}</h3><p>${escapeHtml(data.query_kind)} query · ${formatSeconds(data.elapsed_seconds)} GPU pass · ${escapeHtml(data.cache_state)} cache</p></div><div class="top-score"><b>${formatScore(data.top_match.score)}</b><small>COSINE</small></div></section>
    <nav class="result-tabs" aria-label="Result view"><button class="active" data-result-tab="rank">RANKING</button><button data-result-tab="visual">VISUAL FIELD</button><button data-result-tab="analysis">ANALYSIS</button></nav>
    <div class="result-view" data-result-view="rank">${rankingMarkup(data.rankings)}</div>
    <div class="result-view" data-result-view="visual" hidden>${visualMarkup(data.rankings)}</div>
    <div class="result-view" data-result-view="analysis" hidden>${analysisMarkup(data)}</div>`;
  bindResultTabs(root);
  root.scrollIntoView({ behavior: "smooth", block: "start" });
}

function renderCompareResults(data) {
  const ring = Math.max(0, Math.min(100, (Number(data.selected_score) + 1) * 50));
  const series = [{ title: `${data.query_kind} → ${data.candidate_kind}`, values: data.scores }];
  $("#compare-results").innerHTML = `
    <section class="alignment-card"><div class="score-ring" style="--score:${ring}"><div><b>${formatScore(data.selected_score)}</b><small>COSINE</small></div></div><div class="alignment-copy"><span>PAIRWISE READOUT</span><h3>${escapeHtml(data.label)}</h3><p>${escapeHtml(data.explanation)}</p><div class="alignment-meta"><span>${escapeHtml(data.query_kind)}</span><span>→</span><span>${escapeHtml(data.candidate_kind)}</span><span>${Number(data.dimension).toLocaleString()}D</span><span>${formatSeconds(data.elapsed_seconds)}</span></div></div></section>
    <div class="compare-viz-grid">${lineChart(data.dimensions, series, "Semantic alignment under compression?")}<div class="fingerprint-pair">${fingerprintChart(data.query_fingerprint, `A · ${data.query_kind}`, data.dimension)}${fingerprintChart(data.candidate_fingerprint, `B · ${data.candidate_kind}`, data.dimension)}</div></div>
    <div style="margin-top:12px">${telemetryMarkup(data)}</div>`;
  $("#compare-results").scrollIntoView({ behavior: "smooth", block: "start" });
}

async function submitSearch() {
  const button = $("#search-button");
  clearObjectUrls();
  state.candidateMedia.forEach((file, index) => state.objectUrls.set(`custom-media-${index + 1}`, URL.createObjectURL(file)));
  const payload = {
    query_text: $("#query-text").value,
    query_image: state.searchImage ? handle_file(state.searchImage) : null,
    query_video: state.searchVideo ? handle_file(state.searchVideo) : null,
    custom_texts: $("#custom-texts").value,
    candidate_media: state.candidateMedia.map((file) => handle_file(file)),
    include_showcase: $("#include-showcase").checked,
    dimension: state.searchDimension,
  };
  button.disabled = true;
  startJob("search", Boolean(state.searchVideo || state.candidateMedia.some((file) => file.type.startsWith("video/")) || payload.include_showcase));
  try {
    const client = await connectClient();
    const result = await client.predict("/search", payload);
    const data = unwrapResult(result);
    if (!data || typeof data !== "object") throw new Error("The server returned an unexpected search result.");
    renderSearchResults(data);
    finishJob();
  } catch (error) {
    failJob(error);
  } finally {
    button.disabled = false;
  }
}

async function submitCompare() {
  const button = $("#compare-button");
  const payload = {
    query_text: $("#pair-query-text").value,
    query_image: state.pairQueryImage ? handle_file(state.pairQueryImage) : null,
    query_video: state.pairQueryVideo ? handle_file(state.pairQueryVideo) : null,
    candidate_text: $("#pair-candidate-text").value,
    candidate_image: state.pairCandidateImage ? handle_file(state.pairCandidateImage) : null,
    candidate_video: state.pairCandidateVideo ? handle_file(state.pairCandidateVideo) : null,
    dimension: state.compareDimension,
  };
  button.disabled = true;
  startJob("compare", Boolean(state.pairQueryVideo || state.pairCandidateVideo));
  try {
    const client = await connectClient();
    const result = await client.predict("/compare", payload);
    const data = unwrapResult(result);
    if (!data || typeof data !== "object") throw new Error("The server returned an unexpected comparison result.");
    renderCompareResults(data);
    finishJob();
  } catch (error) {
    failJob(error);
  } finally {
    button.disabled = false;
  }
}

async function loadAssetAsFile(url, filename, type) {
  const response = await fetch(url);
  if (!response.ok) throw new Error(`Could not load the curated asset ${filename}.`);
  return new File([await response.blob()], filename, { type });
}

async function runExample(name) {
  const example = examples[name];
  switchMode("search");
  $("#query-text").value = example.text;
  $("#custom-texts").value = example.custom ?? "";
  $("#include-showcase").checked = example.include ?? true;
  state.searchImage = null;
  state.searchVideo = null;
  fileLabel($("#query-image-drop"), null);
  fileLabel($("#query-video-drop"), null);
  $("#query-image").value = "";
  $("#query-video").value = "";
  setDimension("search", example.dimension);
  if (example.asset) {
    try {
      const file = await loadAssetAsFile(example.asset, example.filename, example.type);
      if (example.type.startsWith("image")) {
        state.searchImage = file;
        fileLabel($("#query-image-drop"), file);
      } else {
        state.searchVideo = file;
        fileLabel($("#query-video-drop"), file);
      }
    } catch (error) {
      failJob(error);
      return;
    }
  }
  setTimeout(submitSearch, 350);
}

function resetSearch() {
  state.searchImage = null;
  state.searchVideo = null;
  state.candidateMedia = [];
  setDimension("search", 1024);
  setTimeout(() => {
    fileLabel($("#query-image-drop"), null);
    fileLabel($("#query-video-drop"), null);
    renderCandidateChips();
  });
}

function bindEvents() {
  $$(".nav-pill").forEach((button) => button.addEventListener("click", () => switchMode(button.dataset.mode)));
  $$("[data-dimension-group]").forEach((group) => $$('button', group).forEach((button) => button.addEventListener("click", () => setDimension(group.dataset.dimensionGroup, button.dataset.value))));
  $$(".expedition").forEach((button) => button.addEventListener("click", () => runExample(button.dataset.example)));

  configureDropzone("#query-image-drop", "#query-image", "searchImage", "searchVideo", "#query-video-drop");
  configureDropzone("#query-video-drop", "#query-video", "searchVideo", "searchImage", "#query-image-drop");
  $("#candidate-media").addEventListener("change", (event) => {
    state.candidateMedia = [...event.target.files].slice(0, 6);
    renderCandidateChips();
  });

  bindPairFile("#pair-query-image", "pairQueryImage", "pairQueryVideo", "#pair-query-file");
  bindPairFile("#pair-query-video", "pairQueryVideo", "pairQueryImage", "#pair-query-file");
  bindPairFile("#pair-candidate-image", "pairCandidateImage", "pairCandidateVideo", "#pair-candidate-file");
  bindPairFile("#pair-candidate-video", "pairCandidateVideo", "pairCandidateImage", "#pair-candidate-file");

  $("#search-form").addEventListener("submit", (event) => { event.preventDefault(); submitSearch(); });
  $("#search-form").addEventListener("reset", resetSearch);
  $("#compare-form").addEventListener("submit", (event) => { event.preventDefault(); submitCompare(); });
  $("#query-text").addEventListener("keydown", (event) => {
    if ((event.metaKey || event.ctrlKey) && event.key === "Enter") {
      event.preventDefault();
      submitSearch();
    }
  });
  $("#job-close").addEventListener("click", () => { $("#job-toast").hidden = true; });
  $("#error-toast button").addEventListener("click", () => { $("#error-toast").hidden = true; });
}

bindEvents();
connectClient().catch(() => {});