let duckdb = null;
let pipeline = null;
let env = null;

const PANEL = document.getElementById("semantic-search-panel");
const TOGGLE = document.getElementById("ssp-toggle");
const CLOSE = document.getElementById("ssp-close");
const FORM = document.getElementById("ssp-form");
const INPUT = document.getElementById("ssp-query");
const TOPK = document.getElementById("ssp-topk");
const SUBMIT = document.getElementById("ssp-submit");
const STATUS = document.getElementById("ssp-status");
const RESULTS = document.getElementById("ssp-results");

let db = null;
let conn = null;
let pca = null;
let encoder = null;
let initPromise = null;
let activeRowId = null;
let highlightedRowIds = [];
const DETAIL_CACHE_LIMIT = 1000;
const detailCache = new Map();

const setStatus = (msg, isError = false) => {
  STATUS.textContent = msg;
  STATUS.classList.toggle("error", !!isError);
};

async function initDuckDB() {
  const bundles = duckdb.getJsDelivrBundles();
  const bundle = await duckdb.selectBundle(bundles);
  const workerUrl = URL.createObjectURL(
    new Blob([`importScripts("${bundle.mainWorker}");`], { type: "text/javascript" }),
  );
  const worker = new Worker(workerUrl);
  const logger = new duckdb.ConsoleLogger(duckdb.LogLevel.WARNING);
  const instance = new duckdb.AsyncDuckDB(logger, worker);
  await instance.instantiate(bundle.mainModule, bundle.pthreadWorker);
  URL.revokeObjectURL(workerUrl);
  await instance.registerFileURL(
    "dataset.parquet",
    new URL("./data/dataset.parquet", window.location.href).toString(),
    duckdb.DuckDBDataProtocol.HTTP,
    false,
  );
  await instance.registerFileURL(
    "details.parquet",
    new URL("./data/details.parquet", window.location.href).toString(),
    duckdb.DuckDBDataProtocol.HTTP,
    false,
  );
  return { db: instance, conn: await instance.connect() };
}

async function loadPCA() {
  const r = await fetch("./data/pca_model.json");
  if (!r.ok) throw new Error(`pca_model.json: HTTP ${r.status}`);
  const m = await r.json();
  return {
    mean: Float32Array.from(m.mean),
    components: m.components.map((row) => Float32Array.from(row)),
    nComponents: m.n_components,
  };
}

async function loadEncoder() {
  return pipeline("feature-extraction", "Xenova/all-MiniLM-L6-v2", {
    quantized: true,
  });
}

async function canUseBrowserCache() {
  if (!globalThis.isSecureContext || !("caches" in globalThis)) return false;
  const cacheName = "nemotron-transformers-cache-test";
  try {
    await caches.open(cacheName);
    await caches.delete(cacheName);
    return true;
  } catch {
    return false;
  }
}

function projectPCA(vec384, pca) {
  const centered = new Float32Array(384);
  for (let i = 0; i < 384; i++) centered[i] = vec384[i] - pca.mean[i];
  const out = new Float32Array(pca.nComponents);
  for (let k = 0; k < pca.nComponents; k++) {
    const row = pca.components[k];
    let s = 0;
    for (let i = 0; i < 384; i++) s += row[i] * centered[i];
    out[k] = s;
  }
  return out;
}

function normalizeVector(vec) {
  let norm = 0;
  for (const v of vec) norm += v * v;
  norm = Math.sqrt(norm);
  if (!norm) return vec;
  for (let i = 0; i < vec.length; i++) vec[i] /= norm;
  return vec;
}

async function encodeQuery(text) {
  const out = await encoder(text, { pooling: "mean", normalize: true });
  return out.data;
}

function toFiniteNumber(value, fallback = 0) {
  const n = typeof value === "number" ? value : Number(value);
  return Number.isFinite(n) ? n : fallback;
}

function getCachedDetail(rowId) {
  const detail = detailCache.get(rowId);
  if (!detail) return null;
  detailCache.delete(rowId);
  detailCache.set(rowId, detail);
  return detail;
}

function cacheDetail(rowId, detail) {
  detailCache.delete(rowId);
  detailCache.set(rowId, detail);
  while (detailCache.size > DETAIL_CACHE_LIMIT) {
    detailCache.delete(detailCache.keys().next().value);
  }
}

function mergeDetails(rows) {
  return rows.map((row) => ({
    ...getCachedDetail(row.row_id),
    ...row,
    loading: !detailCache.has(row.row_id),
  }));
}

async function rankSearch(query, topK, onPhase = () => {}) {
  const t0 = performance.now();
  onPhase("Encoding query…");
  const raw = await encodeQuery(query);
  const vec384 = raw instanceof Float32Array ? raw : Float32Array.from(raw);
  const vec64 = normalizeVector(projectPCA(vec384, pca));
  const tEnc = performance.now() - t0;

  const literal = "[" + Array.from(vec64).map((v) => v.toFixed(7)).join(",") + "]::FLOAT[]";
  const rankSql = `
    SELECT row_id,
           CAST(list_inner_product(emb, ${literal}) AS DOUBLE) AS score
    FROM 'dataset.parquet'
    ORDER BY score DESC
    LIMIT ${topK | 0}
  `;
  const t1 = performance.now();
  onPhase("Ranking vectors…");
  const ranked = await conn.query(rankSql);
  const tRank = performance.now() - t1;
  const topRows = ranked.toArray().map((r, i) => {
    const row = r.toJSON();
    const rowId = toFiniteNumber(row.row_id, NaN);
    if (!Number.isFinite(rowId)) throw new Error(`Invalid row_id: ${row.row_id}`);
    return {
      row_id: Math.trunc(rowId),
      score: toFiniteNumber(row.score),
      ord: i,
    };
  });
  return { rows: topRows, tEnc, tRank };
}

async function hydrateDetails(rows, onPhase = () => {}) {
  const missing = rows.filter((row) => !detailCache.has(row.row_id));
  if (!missing.length) return { rows: mergeDetails(rows), tDetails: 0, fetched: 0 };

  const values = missing
    .map((r) => `(${r.ord | 0}, ${r.row_id})`)
    .join(",");
  const rowIds = missing.map((r) => r.row_id).join(",");
  const detailSql = `
    WITH missing(ord, row_id) AS (VALUES ${values})
    SELECT d.row_id, d.text_preview AS text, d.dataset, d.license
    FROM missing
    JOIN (
      SELECT row_id, text_preview, dataset, license
      FROM 'details.parquet'
      WHERE row_id IN (${rowIds})
    ) AS d ON d.row_id = missing.row_id
    ORDER BY missing.ord
  `;
  const t2 = performance.now();
  onPhase(`Loading ${missing.length} result details…`);
  const detail = await conn.query(detailSql);
  const tDetails = performance.now() - t2;
  for (const r of detail.toArray()) {
    const row = r.toJSON();
    const rowId = Math.trunc(toFiniteNumber(row.row_id, NaN));
    if (Number.isFinite(rowId)) {
      cacheDetail(rowId, {
        row_id: rowId,
        text: row.text,
        dataset: row.dataset,
        license: row.license,
      });
    }
  }
  return { rows: mergeDetails(rows), tDetails, fetched: missing.length };
}

function escapeHtml(s) {
  return String(s ?? "").replace(/[&<>"']/g, (c) => ({
    "&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;", "'": "&#39;",
  })[c]);
}

function normalizeRowId(value) {
  const id = Math.trunc(toFiniteNumber(value, NaN));
  return Number.isFinite(id) ? id : null;
}

function getRowIds(rows) {
  const ids = rows
    .map((row) => normalizeRowId(row.row_id))
    .filter((id) => id !== null);
  return [...new Set(ids)];
}

function syncAtlasSemanticHighlights() {
  if (highlightedRowIds.length) {
    window.__nemotronAtlas?.setSemanticHighlights?.(highlightedRowIds);
  } else {
    window.__nemotronAtlas?.clearSemanticHighlights?.();
  }
}

function setAtlasHighlights(rows) {
  highlightedRowIds = getRowIds(rows);
  syncAtlasSemanticHighlights();
}

function syncAtlasFocus() {
  if (activeRowId !== null) {
    window.__nemotronAtlas?.focusSemanticResult?.(activeRowId);
  } else {
    window.__nemotronAtlas?.clearSemanticFocus?.();
  }
}

function clearAtlasFocus() {
  activeRowId = null;
  window.__nemotronAtlas?.clearSemanticFocus?.();
  RESULTS.querySelectorAll("li.active").forEach((el) => el.classList.remove("active"));
}

function clearAtlasHighlight() {
  activeRowId = null;
  highlightedRowIds = [];
  window.__nemotronAtlas?.clearSemanticHighlights?.();
  window.__nemotronAtlas?.clearSemanticFocus?.();
  RESULTS.querySelectorAll("li.active").forEach((el) => el.classList.remove("active"));
}

function activateResult(li, rowId) {
  const nextRowId = normalizeRowId(rowId);
  if (nextRowId !== null && nextRowId === activeRowId) {
    clearAtlasFocus();
    return;
  }
  activeRowId = nextRowId;
  RESULTS.querySelectorAll("li.active").forEach((el) => el.classList.remove("active"));
  if (activeRowId !== null) li.classList.add("active");
  syncAtlasFocus();
}

function renderResults(rows) {
  RESULTS.innerHTML = "";
  if (!rows.length) {
    RESULTS.innerHTML = `<li style="border:none;color:#6b7280;">No results.</li>`;
    return;
  }
  for (const r of rows) {
    const rowId = normalizeRowId(r.row_id);
    const li = document.createElement("li");
    li.classList.toggle("loading", !!r.loading);
    li.classList.toggle("active", rowId !== null && rowId === activeRowId);
    li.dataset.rowId = String(rowId ?? r.row_id);
    const score = toFiniteNumber(r.score);
    const meta = [
      r.dataset
        ? `<span class="ssp-badge">${escapeHtml(r.dataset)}</span>`
        : `<span class="ssp-badge pending">row ${escapeHtml(r.row_id)}</span>`,
      r.license ? `<span class="ssp-badge license">${escapeHtml(r.license)}</span>` : "",
      `<span class="ssp-badge score">${score.toFixed(3)}</span>`,
    ].filter(Boolean).join("");
    const text = r.text ?? "Loading result details…";
    li.innerHTML = `
      <div class="ssp-result-meta">${meta}</div>
      <div class="ssp-result-text${r.loading ? " pending" : ""}">${escapeHtml(text)}</div>
    `;
    li.addEventListener("click", () => {
      li.classList.toggle("expanded");
      activateResult(li, rowId);
    });
    RESULTS.appendChild(li);
  }
}

async function init() {
  if (encoder && conn && pca) return;
  setStatus("Loading dependencies…");
  duckdb = await import("https://cdn.jsdelivr.net/npm/@duckdb/duckdb-wasm@1.29.0/+esm");
  ({ pipeline, env } = await import("https://cdn.jsdelivr.net/npm/@xenova/transformers@2.17.2"));
  env.allowLocalModels = false;
  env.useBrowserCache = await canUseBrowserCache();
  setStatus("Initializing DuckDB…");
  ({ db, conn } = await initDuckDB());
  setStatus("Loading PCA model…");
  pca = await loadPCA();
  setStatus(
    env.useBrowserCache
      ? "Loading MiniLM encoder (~22 MB, cached after first load)…"
      : "Loading MiniLM encoder (~22 MB)…",
  );
  encoder = await loadEncoder();
  setStatus("Warming up…");
  await encodeQuery("warmup");
  setStatus("Ready.");
  SUBMIT.disabled = false;
  SUBMIT.textContent = "Search";
}

function ensureInit() {
  if (!initPromise) {
    initPromise = init().catch((err) => {
      console.error(err);
      setStatus(`Init failed: ${err.message || err}`, true);
      SUBMIT.disabled = true;
      SUBMIT.textContent = "Unavailable";
      initPromise = null;
      throw err;
    });
  }
  return initPromise;
}

TOGGLE.addEventListener("click", () => {
  PANEL.setAttribute("data-state", "open");
  ensureInit().catch(() => {});
});

CLOSE.addEventListener("click", () => {
  PANEL.setAttribute("data-state", "collapsed");
  clearAtlasHighlight();
});

window.addEventListener("nemotron-atlas-ready", () => {
  syncAtlasSemanticHighlights();
  if (activeRowId !== null) syncAtlasFocus();
});

if (PANEL.getAttribute("data-state") === "open") {
  ensureInit().catch(() => {});
}

FORM.addEventListener("submit", async (e) => {
  e.preventDefault();
  const q = INPUT.value.trim();
  if (!q) return;
  SUBMIT.disabled = true;
  SUBMIT.textContent = "Searching…";
  try {
    await ensureInit();
    clearAtlasHighlight();
    const topK = Math.max(1, Math.min(200, Number(TOPK.value) || 25));
    const ranked = await rankSearch(q, topK, setStatus);
    setAtlasHighlights(ranked.rows);
    renderResults(mergeDetails(ranked.rows));
    const details = await hydrateDetails(ranked.rows, setStatus);
    renderResults(details.rows);
    setStatus(
      `${details.rows.length} results · encode ${ranked.tEnc.toFixed(0)} ms · ` +
      `rank ${ranked.tRank.toFixed(0)} ms · details ${details.tDetails.toFixed(0)} ms`,
    );
  } catch (err) {
    console.error(err);
    setStatus(`Error: ${err.message || err}`, true);
  } finally {
    SUBMIT.disabled = false;
    SUBMIT.textContent = "Search";
  }
});