import * as ort from "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.27.0/dist/ort.webgpu.bundle.min.mjs";
import { AutoTokenizer, env } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.0.2";
import { extractFinalJson, filterVerbatim, buildPrompt, buildUsageImpactPrompt, sentenceContaining } from "./lib/provenance.js";

const HF = (repo, file) => `https://huggingface.co/${repo}/resolve/main/${file}`;
const SPAN_REPO = "rafmacalaba/gliner_datause_extended-onnx";
const SPAN2_REPO = "rafmacalaba/gliner2_datause-onnx";
const MULTITASK_REPO = "rafmacalaba/lfm2.5-350M-datause-multitask-onnx";
const SPAN_LABELS = ["NAMED_DATA", "DESCRIPTIVE_DATA", "VAGUE_DATA"];
const MAX_WIDTH = 12;
const MAX_WIDTH2 = 8;
const MAX_COUNT2 = 20;
let spanThreshold = 0.6;

// Local ONNX files (serve the repo root, not just web/, so ../onnx is reachable).
const LOCAL = {
  span: "../onnx/gliner_datause_extended_fp16.onnx",
  span2: "../onnx/gliner2_datause_fp16_v2.onnx",
  multitaskOnnx: "../onnx/lfm2_multitask_fp16_v2.onnx",
  multitaskData: "../onnx/lfm2_multitask_fp16_v2.onnx_data",
};

env.allowRemoteModels = true;
env.useBrowserCache = true;

const status = document.getElementById("status");
const runBtn = document.getElementById("run");
const results = document.getElementById("results");
const resultsPretty = results.querySelector(".result-pretty");
const resultsJson = results.querySelector(".result-json");
const resultsTrace = results.querySelector(".result-trace");
const segOpts = document.querySelectorAll(".seg-opt");
const segJson = document.querySelector('.seg-opt[data-mode="json"]');
const segTrace = document.querySelector('.seg-opt[data-mode="trace"]');
const logEl = document.getElementById("log");

let lastResult = null;
const statusRule = document.getElementById("status-rule-fill");
const statusPct = document.getElementById("status-pct");
const modelStatusText = document.getElementById("model-status-text");
const textEl = document.getElementById("text");

const SAMPLES = [
  "To assess the distributional impact of agricultural interventions, we draw on several complementary data sources that jointly cover the period of interest. The first is the Demographic and Health Survey (DHS), collected by EuStatistics in 2021, which records health and nutrition outcomes for a nationally representative sample of households across all administrative regions. The second is the Ghana Living Standards Survey (GLSS), administered by the Ghana Statistical Service in 2017, which provides detailed consumption, expenditure, and income modules at the household level. Because both surveys report household consumption using broadly comparable definitions, we are able to pool the two samples and estimate a common poverty line across years and districts. We further merge in the World Bank's World Development Indicators (WDI) to control for district-level economic conditions, population density, and access to infrastructure. Following standard practice, we restrict the estimation sample to rural households and exclude observations with implausible consumption values, and we cluster standard errors at the district level throughout. All monetary variables are expressed in constant local currency units to ensure comparability across survey rounds.",
  "Poverty rates in this paper are estimated using the National Family Health Survey (NFHS-3) for India, a large, nationally representative household survey fielded across all states and union territories. We restrict the estimation sample to ever-married women aged 15 to 49 and cluster standard errors at the district level to account for within-district correlation in household outcomes. The survey records a rich set of household characteristics, including caste, religion, asset ownership, and household size, which we use as controls throughout the analysis. Following the existing literature, we construct a household wealth index from the reported asset holdings using principal components analysis, and we measure poverty with a consumption-based welfare aggregate deflated to a common base year. To guard against measurement error, we exclude observations with missing values on any of the core variables and winsorize the dependent variable at the first and ninety-ninth percentiles. All specifications include state fixed effects and a full set of survey-month fixed effects, and we report robust standard errors clustered at the primary sampling unit level.",
  "This paper estimates the returns to secondary education using a within-family identification strategy that compares the labor-market outcomes of siblings who differ in their educational attainment. The central empirical challenge is that ability and family background jointly determine both schooling and later earnings, so naive comparisons of more- and less-educated individuals confound the causal effect of schooling with pre-existing differences in talent and resources. By restricting attention to brothers and sisters raised in the same household, we remove all time-invariant household characteristics that would otherwise bias the estimates. Our main finding is that an additional year of secondary schooling raises hourly earnings substantially, and the effect is larger for women than for men. We show that this result is robust to a wide range of alternative specifications, including controls for birth order, age spacing, and gender composition, and to alternative definitions of the outcome measure. We also document that the returns to schooling vary with local labor-market conditions, being higher in urban areas and in regions with greater industrial concentration. Finally, we provide evidence that the estimated effects are not driven by selective migration or by differential measurement error in self-reported earnings.",
  "Considerable provincial variation in insurance coverage is evident in the H8BS. Two of the surveys contain self-assessed health (in the GSCF, the assessment was actually done by the woman of the household), which at least in industrialized countries has been found to be a good predictor of mortality. We measured schooling by years of schooling in the CHNS and the GSCF, but in the H8BS have left education coded as levels of attainment. The basic CHNS results are reported in Table 4, the results for the Gansu panel in Table 6, and the results for the Health VIII baseline survey in Table 7.",
  "One of the main advantages of the data from the EAC-I is that plot size is measured by GPS. This avoids measurement errors that lead to biased results. We aggregate crop yields at the household level. We convert production to kilograms using World Bank conversion factors. Interestingly, our data indicate that crop yields vary according to the extent to which polygyny is practiced in the commune.",
  "Data are from the 2015-16 National Family Health Survey (NFHS) of India. Panels A-C present 2SLS estimates as specified by equations 1 and 2 and using cross-sectional data. Panel D presents 2SLS estimates using the retrospective woman-year panel data. Column (1) shows estimates for the full sample, while columns (2) and (3) restrict the samples to those living in rural and urban areas, respectively.",
  "A second contribution of this paper is to show the importance of decomposing the effects with respect to distance from the mines. Given the spatial heterogeneity of the results, we explore the effects in an individual-level, difference-in-differences analysis by using spatial lag models to allow for nonlinear effects with distance from mine. We also allow for spillovers across districts, in a district-level analysis. We use two complementary geocoded household data sets to analyze outcomes in Ghana: the Demographic and Health Survey (DHS) and the Ghana Living Standard Survey (GLSS), which provide information on a wide range of welfare outcomes.",
  "We use geospatial data to test the validity of the local continuity assumption around the boundaries of guerrilla-controlled areas. Elevation was obtained from NASA's Shuttle Radar Topography Mission (SRTM). Information on surface water bodies comes from the MERIT Hydro dataset. Agro-climatic yield rasters with a spatial resolution of five arc minutes (nine km) come from the Global Agro-Ecological Zones (GAEZ) project. For all yields, we are using the 30-year average beginning in 1961.",
  "The average energy spending share observed in the HBS 2019 in Bulgaria is 13.1 percent. The largest energy spending component is on electricity (59.3 percent of all energy expenditure), followed by other solid fuels (27.5 percent) and heat energy (9.3 percent). One stream of literature defines the energy poor as all households that spend more than 10 percent of their total expenditure on energy.",
  "In the COVID-19 NLPS round 4 conducted in August 2020, households were asked to compare their current income from various sources to their income in August 2019, and to indicate whether these incomes had increased, decreased, or remained the same.",
  "Natural amenities were calculated using several GIS layers from the GHS layers. Water bodies were used to calculate dummy variables for location near the sea (coast). We used the 250 meters resolution raster MODIS Water Mask (Carroll et al., 2009) for this purpose. Finally, we used the Lloyd's lists of maritime and fluvial ports to account for the presence of ports.",
  "This annex summarizes how the individual indicators from the WGI data sources are assigned to the six governance dimensions corresponding to the six aggregate WGI indicators in the most recent year in which they appear in the WGI dataset. For most data sources, this is 2023 (the most recent year available in the 2024 WGI update).",
  "The application of this fiscal incidence framework requires both microdata and macrodata. The microdata come from the second wave of the National Survey on Household Living Conditions and Agriculture (ECVMA), and the state budget and national account data provide the macrodata. The ECVMA is a survey analogous to the Living Standards Measurement Study-Integrated Survey on Agriculture (LSMS-ISA), a World Bank Group program designed to collect nationally representative data at the household and individual levels on consumption, agriculture, and other social indicators. The second wave of the panel survey in Niger was implemented in 2014, covering urban and rural areas in all regions of the country. The sample was chosen through a random two-stage process, using the 2001 General Census of Population and Housing as the base for the first stage.",
  "In order to proxy for the origin characteristics, a set of economic variables aggregated at the province (and urban/rural) level is used. Average per capita monthly household expenditures and average per capita land areas owned by households are computed by province (and urban/rural) from the VHLSS 2004 in order to capture living standards and the opportunities for agricultural, forestry, and fishery activities. To proxy for the chance of obtaining a job in the state sector, the proportion of state workers in total employment at each sending area is calculated from the VHLSS 2004. In terms of the proportion of workers employed by foreign firms, the number of workers in the latter sector for each location is taken from the Enterprise Survey data (GSO) and divided by total employment.",
  "Following the survey identification, we proceeded to review the indicators in the identified surveys that were related to norms and female labor force participation. Building on the literature on agency and empowerment, we identified a set of indicators and variables repeatedly appearing in the literature as norms-related indicators. The surveys considered included the Demographic and Health Survey (DHS), World Values Survey (WVS), Multiple Indicator Cluster Surveys (MICS), Gallup Poll, and Gender and Adolescence: Global Evidence (GAGE). WVS, though not available for all countries in the region, was selected due to the availability of gender attitude variables that could serve as a proxy for relevant norms. Upon careful review, the latter three were dropped due to limited data availability and duplication of variables already available in the DHS.",
  "According to IOM's Integrated Location Assessment 4, the rate of physical violence, threats, mistrust and returnees feeling unwelcome decreased within 12 months, but still remained. Context indicators such as blocked return, social cohesion and movement restriction remained similar, according to IOM's Return Index 12. Context indicators in Yathrib showed marked improvement, with threats between tribal groups, blocked IDP return, and returnees feeling unwelcome reported before the LPA was signed and absent three months post-signature.",
  "The SEIS 2025 dataset and its reports were published in December 2025 on the UNHCR Microdata Library. A year later, its uptake is broad. Across agencies, the Ukraine Refugee Response Plan (RRP) relied heavily on SEIS, with 200 citations confirmed and roughly 100 more potentially. In the private sector, Deloitte quoted SEIS in its European Labour Market Outlook and its CEE Ukrainian-refugee integration study, both of which draw on SEIS findings. Among governments, Poland uses SEIS to update its 2026 refugee integration strategy and adjust labour-market access provisions, while Romania draws on SEIS for its own integration planning. In research and academia, SEIS microdata pulled from the UNHCR Microdata Library has been used by several universities and cited in peer-reviewed studies on refugee labour-market outcomes and education access. Media coverage of SEIS has also been widespread.",
];

// --- status logging + progress -------------------------------------------------
function log(msg, cls = "") {
  if (!logEl) return;
  const line = document.createElement("div");
  if (cls) line.className = cls;
  line.textContent = `[${new Date().toLocaleTimeString()}] ${msg}`;
  logEl.appendChild(line);
  logEl.scrollTop = logEl.scrollHeight;
}

function showProgress(frac, label) {
  modelStatusText.textContent = label;
  const pctEl = document.getElementById("loader-pct");
  const ruleEl = document.getElementById("loader-rule-fill");
  const truthEl = document.getElementById("loader-truth");
  const loaderEl = document.getElementById("loader");
  const fillRect = document.getElementById("cat-fill-rect");
  const meniscus = document.getElementById("cat-meniscus");
  const pct = Math.round(frac * 100);
  if (pctEl) pctEl.textContent = pct;
  if (ruleEl) ruleEl.style.width = `${pct}%`;
  if (truthEl) truthEl.textContent = label;
  if (statusRule) statusRule.style.width = `${pct}%`;
  if (statusPct) statusPct.textContent = `${pct}%`;
  if (loaderEl) loaderEl.dataset.state = pct < 50 ? "sleeping" : pct < 90 ? "alert" : "excited";

  if (fillRect) {
    const fillY = Math.max(0, Math.min(320, 320 - frac * 320));
    const fillH = Math.max(0, Math.min(320, frac * 320));
    fillRect.setAttribute("y", fillY);
    fillRect.setAttribute("height", fillH);
    if (meniscus) {
      meniscus.setAttribute("y1", fillY);
      meniscus.setAttribute("y2", fillY);
      meniscus.style.opacity = frac > 0.01 && frac < 0.99 ? "1" : "0.4";
    }
  }
}

function setLoaderStatus(msg) {
  const el = document.getElementById("loader-status");
  if (el) el.textContent = msg;
}

function hideLoader() {
  document.getElementById("loader")?.classList.add("hidden");
}

function enableLoaderButton() {
  const btn = document.getElementById("loader-cta");
  if (btn) btn.disabled = false;
}

function setStatus(msg) { 
  status.textContent = msg;
  const pill = document.getElementById("status-pill");
  if (!pill) return;
  const t = /error|unavailable|not found|rejected|failed/i.test(msg) ? "err"
          : /ready|done/i.test(msg) ? "ready"
          : "busy";
  pill.className = "status-pill " + t;
}

// --- Browser model cache --------------------------------------------------------
// Primary: Cache API (same backend transformers.js uses, proven in HF Spaces).
// Fallback: OPFS, used when the Cache API is unavailable or rejects a write.
const CACHE_NAME = "datause-models";
const CACHE_DIR = "datause-models";

function hashKey(str) {
  let h = 0x811c9dc5;
  for (let i = 0; i < str.length; i++) {
    h ^= str.charCodeAt(i);
    h = Math.imul(h, 0x01000193);
  }
  return (h >>> 0).toString(16).padStart(8, "0");
}

// Cache API backend ----------------------------------------------------------
async function readCacheApi(key) {
  if (typeof caches === "undefined") return null;
  try {
    const cache = await caches.open(CACHE_NAME);
    const resp = await cache.match(key);
    if (!resp) return null;
    return await resp.arrayBuffer();
  } catch (e) {
    log(`cache read failed (Cache API): ${e?.name || e}`, "warn");
    return null;
  }
}

async function writeCacheApi(key, buffer) {
  if (typeof caches === "undefined") return false;
  try {
    const cache = await caches.open(CACHE_NAME);
    await cache.put(key, new Response(buffer));
    return true;
  } catch (e) {
    log(`cache write failed (Cache API): ${e?.name || e}`, "warn");
    return false;
  }
}

// OPFS backend ---------------------------------------------------------------
async function opfsCacheDir() {
  const root = await navigator.storage.getDirectory();
  return root.getDirectoryHandle(CACHE_DIR, { create: true });
}

async function readOpfs(key) {
  if (!navigator.storage?.getDirectory) return null;
  try {
    const dir = await opfsCacheDir();
    const name = hashKey(key);
    await dir.getFileHandle(name + ".ok"); // marker proves a complete prior write
    const handle = await dir.getFileHandle(name);
    const file = await handle.getFile();
    return await file.arrayBuffer();
  } catch (e) {
    log(`cache read failed (OPFS): ${e?.name || e}`, "warn");
    return null;
  }
}

async function writeOpfs(key, buffer) {
  if (!navigator.storage?.getDirectory) return false;
  try {
    const dir = await opfsCacheDir();
    const name = hashKey(key);
    try { await dir.removeEntry(name + ".ok"); } catch { /* not present yet */ }
    const handle = await dir.getFileHandle(name, { create: true });
    const writable = await handle.createWritable();
    await writable.write(buffer);
    await writable.close();
    await dir.getFileHandle(name + ".ok", { create: true });
    return true;
  } catch (e) {
    log(`cache write failed (OPFS): ${e?.name || e}`, "warn");
    return false;
  }
}

// Unified entry points -------------------------------------------------------
async function cacheGet(key) {
  const viaCacheApi = await readCacheApi(key);
  if (viaCacheApi) return viaCacheApi;
  return readOpfs(key);
}

async function cacheSet(key, buffer) {
  const ok = await writeCacheApi(key, buffer);
  if (!ok) await writeOpfs(key, buffer);
}

// --- download with progress + cache -------------------------------------------
async function loadFile(url, label) {
  const cached = await cacheGet(url);
  if (cached) { log(`cached: ${label} (${(cached.byteLength / 1e6).toFixed(1)} MB)`, "ok"); return cached; }

  log(`downloading ${label} …`);
  setLoaderStatus("Downloading");
  const resp = await fetch(url);
  if (!resp.ok) throw new Error(`${resp.status} for ${url}`);
  const total = Number(resp.headers.get("content-length") || 0);
  const reader = resp.body.getReader();
  const chunks = [];
  let received = 0;
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    received += value.length;
    const mb = (received / 1e6).toFixed(1);
    if (total) {
      showProgress(received / total, `${label} ${mb} / ${(total / 1e6).toFixed(1)} MB`);
      setStatus(`Downloading ${label}… ${mb} / ${(total / 1e6).toFixed(1)} MB`);
    } else {
      showProgress(0.5, `${label} ${mb} MB`);
      setStatus(`Downloading ${label}… ${mb} MB`);
    }
  }
  const buffer = new Uint8Array(received);
  let off = 0;
  for (const c of chunks) { buffer.set(c, off); off += c.length; }
  log(`downloaded ${label} (${(received / 1e6).toFixed(1)} MB)`, "ok");
  showProgress(1, label);
  cacheSet(url, buffer.buffer)
    .then(() => log(`cached ${label}`, "ok"));
  return buffer.buffer;
}

// --- LFM2.5 generation ----------------------------------------------------------
const HIDDEN = 1024, NUM_KV_HEADS = 8, HEAD_DIM = 64;

async function generateRaw(session, tokenizer, prompt, onToken, stopWhen) {
  const rendered = `<|startoftext|><|im_start|>user\n${prompt}<|im_end|>\n<|im_start|>assistant\n`;
  const ids = tokenizer.encode(rendered, { add_special_tokens: false });

  const cache = {};
  for (const name of session.inputNames) {
    if (name.startsWith("past_conv")) cache[name] = new ort.Tensor("float16", new Uint16Array(HIDDEN * 3), [1, HIDDEN, 3]);
    else if (name.startsWith("past_key_values")) cache[name] = new ort.Tensor("float16", new Uint16Array(0), [1, NUM_KV_HEADS, 0, HEAD_DIM]);
  }

  const eos = tokenizer.eos_token_id;
  const gen = [];
  let cur = ids;
  let lastText = "";
  for (let step = 0; step < 256; step++) {
    const feed = {
      input_ids: new ort.Tensor("int64", new BigInt64Array(cur.map(BigInt)), [1, cur.length]),
      attention_mask: new ort.Tensor("int64", new BigInt64Array(ids.length + gen.length).fill(1n), [1, ids.length + gen.length]),
      num_logits_to_keep: new ort.Tensor("int64", new BigInt64Array([1n]), []),
      ...cache,
    };
    const out = await session.run(feed);
    const logits = out.logits.data;
    const V = out.logits.dims[2];
    const start = (out.logits.dims[1] - 1) * V;
    let best = 0;
    for (let i = 1; i < V; i++) if (logits[start + i] > logits[start + best]) best = i;
    if (best === eos) break;
    gen.push(best);
    lastText = tokenizer.decode(gen, { skip_special_tokens: true });
    if (onToken) onToken(lastText);
    if (stopWhen && stopWhen(lastText)) break;
    for (const [name, t] of Object.entries(out)) {
      if (name.startsWith("present_conv")) cache["past_conv." + name.slice("present_conv.".length)] = t;
      else if (name.startsWith("present.")) cache["past_key_values." + name.slice("present.".length)] = t;
    }
    cur = [best];
  }
  return lastText || tokenizer.decode(gen, { skip_special_tokens: true });
}

// True once a complete JSON object has been emitted (after any <think> block).
function jsonComplete(text) {
  if (text.includes("<think>") && !text.includes("</think>")) return false;
  return extractFinalJson(text) !== null;
}

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

// Render a parsed JSON value as HTML: keys navy-bold (span.k), values plain.
function jsonToHtml(value) {
  if (value === null || value === undefined) return "null";
  const t = typeof value;
  if (t === "string") return escapeHtml(JSON.stringify(value));
  if (t === "number" || t === "boolean") return String(value);
  if (Array.isArray(value)) return "[" + value.map(jsonToHtml).join(", ") + "]";
  if (t === "object") {
    return "{ " + Object.entries(value)
      .map(([k, v]) => '<span class="k">' + escapeHtml(JSON.stringify(k)) + "</span>: " + jsonToHtml(v))
      .join(", ") + " }";
  }
  return escapeHtml(String(value));
}

// --- Pretty tree rendering ------------------------------------------------------
const PROV_KEYS = ["acronym", "producer", "year", "geography"];
const UI_KEYS = ["data_type", "usage_action", "impact_label", "usage_summary"];

function leafHtml(k, v) {
  return `<div class="tree-leaf"><span class="k">${escapeHtml(k)}</span><span class="arrow">→</span><span class="v">${escapeHtml(String(v))}</span></div>`;
}

function leavesHtml(obj, keys) {
  if (!obj) return '<div class="tree-none">none</div>';
  const out = [];
  for (const k of keys) {
    const v = obj[k];
    if (v === undefined || v === null || v === "") continue;
    out.push(leafHtml(k, v));
  }
  return out.length ? out.join("") : '<div class="tree-none">none</div>';
}

function confidenceChip(score) {
  if (typeof score !== "number" || Number.isNaN(score)) return "";
  const p = 1 / (1 + Math.exp(-score));
  const pct = Math.round(p * 100);
  const level = p >= 0.85 ? "c-high" : p >= 0.7 ? "c-mid" : "c-low";
  return `<span class="tree-conf ${level}" title="confidence ${pct}%">${pct}%</span>`;
}

function mentionTreeHtml(m) {
  const badge = `<span class="badge badge-${m.label}">${escapeHtml(m.label)}</span>`;
  const name = `<span class="tree-name">${escapeHtml(m.text)}</span>`;
  const conf = confidenceChip(m.score);
  if (m.label === "VAGUE_DATA") {
    return `<div class="tree"><div class="tree-root">${badge}${name}${conf}</div><div class="tree-none">no attributes</div></div>`;
  }
  return (
    `<div class="tree">` +
    `<div class="tree-root">${badge}${name}${conf}<span class="tree-ms">…</span></div>` +
    `<div class="tree-body">` +
    `<div class="tree-slot" data-slot="prov"><div class="tree-stream streaming">attributing…</div></div>` +
    `<div class="tree-group"><div class="tree-label">usage/impact</div><div class="tree-slot" data-slot="ui"></div></div>` +
    `</div></div>`
  );
}

async function runProvenance(session, tokenizer, text, mention, onToken) {
  const raw = await generateRaw(session, tokenizer, buildPrompt(text, mention), onToken, jsonComplete);
  return filterVerbatim(extractFinalJson(raw), text);
}

async function runUsageImpact(session, tokenizer, text, mention, onToken) {
  const usage = sentenceContaining(text, mention);
  const raw = await generateRaw(session, tokenizer, buildUsageImpactPrompt(mention, usage, text), onToken, jsonComplete);
  return extractFinalJson(raw) ?? {};
}

// --- GLiNER span detection ------------------------------------------------------
function prepareSpanIdx(numWords, maxWidth) {
  const starts = [], ends = [], mask = [];
  for (let s = 0; s < numWords; s++) {
    for (let w = 0; w < maxWidth; w++) {
      const e = s + w;
      starts.push(s);
      ends.push(e);
      mask.push(e < numWords ? 1 : 0);
    }
  }
  return { starts, ends, mask };
}

function tokenizeWords(text) {
  // Approximate spaCy `blank('en')` word splitting: whitespace + punctuation.
  const out = [];
  for (const chunk of text.split(/\s+/)) {
    if (!chunk) continue;
    const parts = chunk.match(/[A-Za-z0-9][A-Za-z0-9'\-]*|[^\s]/g);
    if (parts) out.push(...parts);
  }
  return out;
}

function buildSpanInput(tokenizer, text) {
  // GLiNER uni-encoder prompt: <<ENT>> label, repeated per label, then <<SEP>>.
  const prompt = ["<<ENT>>", "NAMED_DATA", "<<ENT>>", "DESCRIPTIVE_DATA",
                  "<<ENT>>", "VAGUE_DATA", "<<SEP>>"];
  const words = tokenizeWords(text);
  const clsId = tokenizer.cls_token_id ?? 1;
  const sepId = tokenizer.sep_token_id ?? 2;

  const ids = [clsId];
  const wordsMask = [0n];
  for (const w of prompt) {
    for (const s of tokenizer.encode(w, { add_special_tokens: false })) {
      ids.push(s); wordsMask.push(0n);
    }
  }
  let wordId = 0;
  for (const w of words) {
    wordId++;
    const sub = tokenizer.encode(w, { add_special_tokens: false });
    for (let i = 0; i < sub.length; i++) {
      ids.push(sub[i]);
      wordsMask.push(i === 0 ? BigInt(wordId) : 0n);
    }
  }
  ids.push(sepId);
  wordsMask.push(0n);

  return {
    ids,
    wordsMask: new BigInt64Array(wordsMask),
    numWords: wordId,
    words,
  };
}

async function runSpan(session, tokenizer, text) {
  const { ids, wordsMask, numWords, words } = buildSpanInput(tokenizer, text);
  if (numWords <= 0) return [];

  const { starts, ends, mask } = prepareSpanIdx(numWords, MAX_WIDTH);
  const spanIdx = [];
  for (let i = 0; i < starts.length; i++) spanIdx.push([starts[i], ends[i]]);

  const feed = {
    input_ids: new ort.Tensor("int64", new BigInt64Array(ids.map(BigInt)), [1, ids.length]),
    attention_mask: new ort.Tensor("int64", new BigInt64Array(ids.length).fill(1n), [1, ids.length]),
    words_mask: new ort.Tensor("int64", wordsMask, [1, ids.length]),
    text_lengths: new ort.Tensor("int64", new BigInt64Array([BigInt(numWords)]), [1, 1]),
    span_idx: new ort.Tensor("int64", new BigInt64Array(spanIdx.flat().map(BigInt)), [1, spanIdx.length, 2]),
    span_mask: new ort.Tensor("int64", new BigInt64Array(mask.map(BigInt)), [1, mask.length]),
  };
  const out = await session.run(feed);
  const logits = out.logits.data;
  const C = out.logits.dims[3];
  const W = out.logits.dims[2];
  const numSpanTotal = out.logits.dims[1] * W;

  const spans = [];
  for (let i = 0; i < numSpanTotal; i++) {
    if (!mask[i]) continue;
    const base = i * C;
    let cls = 0;
    for (let c = 1; c < C; c++) if (logits[base + c] > logits[base + cls]) cls = c;
    spans.push({ start: starts[i], end: ends[i], cls, score: logits[base + cls] });
  }
  spans.sort((a, b) => b.score - a.score);
  const kept = [];
  const used = new Set();
  for (const sp of spans) {
    if (sp.score < spanThreshold) continue;
    let overlaps = false;
    for (let i = sp.start; i <= sp.end; i++) if (used.has(i)) { overlaps = true; break; }
    if (overlaps) continue;
    for (let i = sp.start; i <= sp.end; i++) used.add(i);
    const label = SPAN_LABELS[sp.cls];
    const txt = words.slice(sp.start, sp.end + 1).join(" ");
    kept.push({ text: txt, label, score: sp.score });
  }
  return kept;
}

async function loadModelFile(localPath, hfUrl, label) {
  // Try local file first (fast, offline); fall back to HuggingFace.
  try {
   const r = await fetch(localPath, { method: "HEAD" });
   if (r.ok) { log(`using local ${label}`); return loadFile(localPath, label); }
  } catch { /* no local server path */ }
  return loadFile(hfUrl, label);
}

// --- GLiNER2 span detection (schema-based) --------------------------------------
function tokenizeWords2(text) {
  // GLiNER2's regex word splitter (lowercased), returning tokens with char offsets.
  const lower = text.toLowerCase();
  const re = /(?:https?:\/\/[^\s]+|www\.[^\s]+)|[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}|@[a-z0-9_]+|\w+(?:[-_]\w+)*|\S/gi;
  const out = [];
  let m;
  while ((m = re.exec(lower)) !== null) {
    out.push({ t: m[0], s: m.index, e: m.index + m[0].length });
  }
  return out;
}

function buildSpan2Input(tokenizer, text) {
  if (text && !/[.!?]$/.test(text)) text += ".";
  const schema = ["(", "[P]", "entities", "(", "[E]", "NAMED_DATA", "[E]",
                  "DESCRIPTIVE_DATA", "[E]", "VAGUE_DATA", ")", ")"];
  const words = tokenizeWords2(text);
  const combined = [...schema, "[SEP_TEXT]", ...words.map((w) => w.t)];

  const ids = [];
  const schemaSpecial = [];
  const textWordIndices = [];
  const startMap = [];
  const endMap = [];

  let pos = 0;
  for (let i = 0; i < combined.length; i++) {
    const w = combined[i];
    const sub = tokenizer.encode(w, { add_special_tokens: false });
    const startPos = pos;
    for (const s of sub) { ids.push(s); pos++; }
    if (i < schema.length && (w === "[P]" || w === "[E]")) schemaSpecial.push(startPos);
    if (i >= schema.length + 1) {
      textWordIndices.push(startPos);
      startMap.push(words[i - schema.length - 1].s);
      endMap.push(words[i - schema.length - 1].e);
    }
  }
  return { ids, schemaSpecial, textWordIndices, startMap, endMap, numWords: words.length };
}

function prepareSpanIdx2(numWords, maxWidth) {
  const starts = [], ends = [], mask = [];
  for (let s = 0; s < numWords; s++) {
    for (let w = 0; w < maxWidth; w++) {
      const e = s + w;
      starts.push(e < numWords ? s : -1);
      ends.push(e < numWords ? e : -1);
      mask.push(e < numWords ? 1 : 0);
    }
  }
  return { starts, ends, mask };
}

async function runSpan2(session, tokenizer, text) {
  const { ids, schemaSpecial, textWordIndices, startMap, endMap, numWords } =
    buildSpan2Input(tokenizer, text);
  if (numWords <= 0) return [];

  const { starts, ends, mask } = prepareSpanIdx2(numWords, MAX_WIDTH2);
  const spanIdx = [];
  for (let i = 0; i < starts.length; i++) {
    const a = starts[i] < 0 ? 0 : starts[i];
    const b = ends[i] < 0 ? 0 : ends[i];
    spanIdx.push([a, b]);
  }

  const feed = {
    input_ids: new ort.Tensor("int64", new BigInt64Array(ids.map(BigInt)), [1, ids.length]),
    attention_mask: new ort.Tensor("int64", new BigInt64Array(ids.length).fill(1n), [1, ids.length]),
    text_word_indices: new ort.Tensor("int64", new BigInt64Array(textWordIndices.map(BigInt)), [1, numWords]),
    schema_special_indices: new ort.Tensor("int64", new BigInt64Array(schemaSpecial.map(BigInt)), [1, schemaSpecial.length]),
    span_idx: new ort.Tensor("int64", new BigInt64Array(spanIdx.flat().map(BigInt)), [1, spanIdx.length, 2]),
  };
  const out = await session.run(feed);
  const scores = out.span_scores.data; // (1, W, K=8, max_count=20, M=3)
  const W = out.span_scores.dims[1];
  const K = out.span_scores.dims[2];
  const MC = out.span_scores.dims[3];
  const M = out.span_scores.dims[4];

  // slot 0 scores: (W, K, M)
  const kept = [];
  for (let mIdx = 0; mIdx < M; mIdx++) {
    const label = SPAN_LABELS[mIdx];
    const cand = [];
    for (let w = 0; w < W; w++) {
      for (let k = 0; k < K; k++) {
        if (!mask[w * K + k]) continue;
        const idx = (((w * K + k) * MC + 0) * M) + mIdx;
        const sc = scores[idx];
        if (sc >= spanThreshold) {
          const start = w;
          const end = w + k; // inclusive word span [w, w+k]
          cand.push({ start, end, score: sc });
        }
      }
    }
    cand.sort((a, b) => b.score - a.score);
    const used = new Set();
    for (const c of cand) {
      let overlap = false;
      for (let i = c.start; i <= c.end; i++) if (used.has(i)) { overlap = true; break; }
      if (overlap) continue;
      for (let i = c.start; i <= c.end; i++) used.add(i);
      const charStart = startMap[c.start];
      const charEnd = endMap[c.end];
      const txt = text.slice(charStart, charEnd).trim();
      if (txt) kept.push({ text: txt, label });
    }
  }
  return kept;
}

// --- init -----------------------------------------------------------------------
async function init() {
  if (!navigator.gpu) { setStatus("WebGPU unavailable"); log("WebGPU unavailable", "warn"); return null; }
  if (navigator.storage?.persist) {
    navigator.storage.persist()
      .then((ok) => log(`persistent storage ${ok ? "granted" : "not granted"}`, ok ? "ok" : "warn"))
      .catch(() => {});
  }
  log(`env: caches=${typeof caches !== "undefined"} opfs=${!!navigator.storage?.getDirectory} embedded=${window.self !== window.top}`, "ok");
  if (navigator.storage?.estimate) {
    navigator.storage.estimate()
      .then(({ usage, quota }) => log(`storage: ${(usage / 1e6).toFixed(1)} MB used / ${(quota / 1e6).toFixed(1)} MB quota`, "ok"))
      .catch(() => {});
  }
  setStatus("Requesting WebGPU adapter…");
  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) { setStatus("WebGPU adapter not found"); log("WebGPU adapter not found", "warn"); return null; }
  ort.env.wasm.numThreads = 1;

  setStatus("Loading span model (GLiNER)…");
  log("loading span model (GLiNER) …");
  const spanBuffer = await loadModelFile(LOCAL.span, HF(SPAN_REPO, "onnx/gliner_datause_extended_fp16.onnx"), "GLiNER span model");
  setStatus("Initializing span model (WebGPU)…");
  const spanSession = await ort.InferenceSession.create(spanBuffer, { executionProviders: ["webgpu"] });
  log("span model ready", "ok");

  setStatus("Loading tokenizers…");
  log("loading tokenizers …");
  const spanTokenizer = await AutoTokenizer.from_pretrained("rafmacalaba/gliner_datause_extended-onnx");
  const multitaskTokenizer = await AutoTokenizer.from_pretrained(MULTITASK_REPO);
  log("tokenizers ready", "ok");

  const multitaskSession = await loadMultitaskSession();

  setStatus("Ready");
  modelStatusText.textContent = "Models ready — cached in browser · WebGPU";
  if (statusRule) statusRule.style.width = "100%";
  if (statusPct) statusPct.textContent = "100%";
  runBtn.disabled = false;
  setLoaderStatus("Ready");
  log("execution provider: WebGPU (all models)", "ok");
  enableLoaderButton();
  return { spanSession, spanTokenizer, multitaskTokenizer, multitaskSession };
}

// --- LFM2.5 multitask attribution ------------------------------------------------
async function createLlmSession(onnxBuffer, dataBuffer, dataPath, label) {
  const extData = [{ path: dataPath, data: dataBuffer }];
  setStatus(`Initializing ${label} (WebGPU)…`);
  try {
    return await ort.InferenceSession.create(onnxBuffer, { executionProviders: ["webgpu"], externalData: extData });
  } catch (e) {
    log(`WebGPU rejected ${label}: ${e?.message || e}`, "warn");
    console.error("webgpu create error:", e);
    setStatus(`WebGPU rejected ${label} — trying wasm (CPU)…`);
    const s = await ort.InferenceSession.create(onnxBuffer, { executionProviders: ["wasm"], externalData: extData });
    log(`${label} loaded on wasm (CPU) fallback`, "warn");
    return s;
  }
}

async function loadMultitaskSession() {
  setStatus("Loading multitask model (LFM2.5)…");
  log("loading multitask model (LFM2.5) …");
  const onnx = await loadModelFile(LOCAL.multitaskOnnx, HF(MULTITASK_REPO, "onnx/lfm2_multitask_fp16_v2.onnx"), "LFM2.5 multitask model");
  const data = await loadModelFile(LOCAL.multitaskData, HF(MULTITASK_REPO, "onnx/lfm2_multitask_fp16_v2.onnx_data"), "LFM2.5 multitask weights");
  const s = await createLlmSession(onnx, data, "lfm2_multitask_fp16_v2.onnx_data", "multitask model");
  log("multitask model ready", "ok");
  return s;
}

const ready = init().catch((e) => {
  const msg = (e && (e.message || e.name)) || String(e) || "unknown error";
  setStatus("Error: " + msg);
  setLoaderStatus("Error");
  log("Error: " + msg, "warn");
  log("details: " + (e?.stack || JSON.stringify(e) || "n/a"), "warn");
  console.error(e);
  enableLoaderButton();
  return null;
});

document.getElementById("loader-cta")?.addEventListener("click", () => {
  hideLoader();
});

runBtn.addEventListener("click", async () => {
  const ctx = await ready;
  if (!ctx) return;
  const text = textEl.value.trim();
  if (!text) return;
  resultsPretty.innerHTML = "";
  resultsJson.innerHTML = "";
  resultsTrace.innerHTML = "";
  lastResult = null;
  [segJson, segTrace].forEach((s) => { if (s) s.disabled = true; });
  runBtn.disabled = true;

  try {
    const tTotal0 = performance.now();

    setStatus("Running span model (GLiNER)…");
    log("extracting data mentions (GLiNER) …");
    const tSpan0 = performance.now();
    const mentions = await runSpan(ctx.spanSession, ctx.spanTokenizer, text);
    const spanMs = performance.now() - tSpan0;
    const named = mentions.filter((m) => m.label !== "VAGUE_DATA");
    log(`span: ${spanMs.toFixed(0)} ms — ${mentions.length} mentions (${named.length} named/descriptive)`, "ok");

    const labelCounts = {};
    for (const m of mentions) labelCounts[m.label] = (labelCounts[m.label] || 0) + 1;
    log(`labels: ${SPAN_LABELS.map(l => `${l}=${labelCounts[l] || 0}`).join(", ")}`, "ok");

    // Attribution: LFM2.5 multitask (provenance + usage/impact).
    const attribLabel = "Multitask";

    // Render all mentions immediately; attribution fills in afterwards.
    const perfDiv = document.createElement("div");
    perfDiv.className = "perf";
    perfDiv.innerHTML =
      `<span class="stat">span <b>${spanMs.toFixed(0)} ms</b></span>` +
      `<span class="stat">mentions <b>${named.length}</b></span>` +
      `<span class="stat">${attribLabel.toLowerCase()} <b>…</b></span>`;
    resultsPretty.prepend(perfDiv);
    const perfAttrib = perfDiv.querySelector(".stat:last-child b");

    const cards = new Map();
    for (const m of mentions) {
      const wrapper = document.createElement("div");
      wrapper.innerHTML = mentionTreeHtml(m);
      const el = wrapper.firstElementChild;
      resultsPretty.appendChild(el);
      cards.set(m, {
        msEl: el.querySelector(".tree-ms"),
        provSlot: el.querySelector('[data-slot="prov"]'),
        uiSlot: el.querySelector('[data-slot="ui"]'),
        attrs: null, ui: null, ms: null,
      });
    }

    // Attribution per mention, streaming tokens into each tree as they arrive.
    let attribTotalMs = 0;
    for (const m of named) {
      setStatus(`Attributing "${m.text}" (${attribLabel})…`);
      const t0 = performance.now();
      const card = cards.get(m);

      const attrs = await runProvenance(ctx.multitaskSession, ctx.multitaskTokenizer, text, m.text, (partial) => {
        if (card && card.provSlot) card.provSlot.innerHTML = `<div class="tree-stream streaming">${escapeHtml(partial)}</div>`;
      });
      if (card && card.provSlot) {
        card.attrs = attrs;
        card.provSlot.innerHTML = leavesHtml(attrs, PROV_KEYS);
      }

      const ui = await runUsageImpact(ctx.multitaskSession, ctx.multitaskTokenizer, text, m.text, (partial) => {
        if (card && card.uiSlot) card.uiSlot.innerHTML = `<div class="tree-stream streaming">${escapeHtml(partial)}</div>`;
      });
      if (card && card.uiSlot) {
        card.ui = ui;
        card.uiSlot.innerHTML = leavesHtml(ui, UI_KEYS);
      }

      const ms = performance.now() - t0;
      attribTotalMs += ms;
      log(`attribution "${m.text}": ${ms.toFixed(0)} ms`);
      if (card) {
        card.ms = ms;
        if (card.msEl) card.msEl.textContent = `${ms.toFixed(0)} ms`;
      }
      perfAttrib.textContent = `${attribTotalMs.toFixed(0)} ms`;
    }

    const totalMs = performance.now() - tTotal0;
    log(`done: total ${totalMs.toFixed(0)} ms · span ${spanMs.toFixed(0)} ms · ${attribLabel.toLowerCase()} ${attribTotalMs.toFixed(0)} ms`, "ok");
    if (!mentions.length) resultsPretty.innerHTML = "<div class=\"empty\">No mentions found.</div>";
    setStatus("Done");

    lastResult = {
      input: { text, length: text.length },
      timings: {
        span_ms: Math.round(spanMs),
        attribution_ms: Math.round(attribTotalMs),
        total_ms: Math.round(totalMs),
      },
      label_counts: labelCounts,
      mentions: mentions.map((m) => {
        const c = cards.get(m) || {};
        return {
          text: m.text,
          label: m.label,
          provenance: c.attrs || null,
          usage_impact: c.ui || null,
          attribution_ms: c.ms != null ? Math.round(c.ms) : null,
        };
      }),
    };
    [segJson, segTrace].forEach((s) => { if (s) s.disabled = false; });
    const activeMode = document.querySelector(".seg-opt.active")?.dataset.mode || "pretty";
    if (activeMode === "json") renderJson();
    else if (activeMode === "trace") renderTrace();
  } catch (e) {
    const msg = (e && (e.message || e.name)) || String(e) || "unknown error";
    setStatus("Error: " + msg);
    log("Error: " + msg, "warn");
    log("details: " + (e?.stack || JSON.stringify(e) || "n/a"), "warn");
    console.error(e);
  } finally {
    runBtn.disabled = false;
  }
});

function colorizeJson(json) {
  return escapeHtml(json)
    .replace(/(&quot;[^&]*?&quot;)(\s*:)/g, '<span class="jk">$1</span>$2')
    .replace(/:\s*(&quot;[^&]*?&quot;)/g, ': <span class="js">$1</span>')
    .replace(/:\s*(-?\d+\.?\d*)/g, ': <span class="jn">$1</span>')
    .replace(/:\s*(true|false|null)\b/g, ': <span class="jb">$1</span>');
}

function cleanResult() {
  if (!lastResult) return null;
  return {
    input: lastResult.input,
    mentions: lastResult.mentions.map((m) => ({
      text: m.text,
      label: m.label,
      provenance: m.provenance,
      usage_impact: m.usage_impact,
    })),
  };
}

function renderJson() {
  if (!lastResult) return;
  const json = JSON.stringify(cleanResult(), null, 2);
  resultsJson.innerHTML = `<pre>${colorizeJson(json)}</pre>`;
}

function renderTrace() {
  if (!lastResult) return;
  const t = lastResult.timings;
  const meta =
    `<div class="json-meta">` +
    `<span>input <b>${lastResult.input.length}</b> chars</span>` +
    `<span>span <b>${t.span_ms} ms</b></span>` +
    `<span>attribution <b>${t.attribution_ms} ms</b></span>` +
    `<span>total <b>${t.total_ms} ms</b></span>` +
    `<span>mentions <b>${lastResult.mentions.length}</b></span>` +
    `</div>`;
  const json = JSON.stringify(lastResult, null, 2);
  resultsTrace.innerHTML = meta + `<pre>${colorizeJson(json)}</pre>`;
}

const PANES = [
  ["pretty", resultsPretty],
  ["json", resultsJson],
  ["trace", resultsTrace],
];

function setMode(mode) {
  if (mode !== "pretty" && !lastResult) return;
  segOpts.forEach((o) => {
    const active = o.dataset.mode === mode;
    o.classList.toggle("active", active);
    o.setAttribute("aria-selected", String(active));
  });
  for (const [m, pane] of PANES) {
    pane.hidden = m !== mode;
    pane.setAttribute("aria-live", m === mode ? "polite" : "off");
  }
  if (mode === "json") renderJson();
  else if (mode === "trace") renderTrace();
}

segOpts.forEach((o) => o.addEventListener("click", () => setMode(o.dataset.mode)));

// Sample category tabs & chips
const sampleSegBtns = document.querySelectorAll(".sample-seg-btn");
const sampleGroups = document.querySelectorAll(".samples-group");
const sampleChips = document.querySelectorAll(".sample-chip");
const shuffleBtn = document.getElementById("sample-shuffle");

function selectSample(btn) {
  if (!btn) return;
  const idx = Number(btn.dataset.sample);
  if (isNaN(idx) || !SAMPLES[idx]) return;
  textEl.value = SAMPLES[idx];
  sampleChips.forEach((c) => c.classList.toggle("active", c === btn));
}

sampleSegBtns.forEach((btn) => {
  btn.addEventListener("click", () => {
    const cat = btn.dataset.category;
    sampleSegBtns.forEach((b) => {
      const active = b === btn;
      b.classList.toggle("active", active);
      b.setAttribute("aria-selected", String(active));
    });
    sampleGroups.forEach((g) => {
      g.hidden = g.dataset.group !== cat;
    });
  });
});

sampleChips.forEach((btn) => {
  btn.addEventListener("click", () => selectSample(btn));
});

shuffleBtn?.addEventListener("click", () => {
  const activeGroup = document.querySelector(".samples-group:not([hidden])");
  if (!activeGroup) return;
  const visibleChips = Array.from(activeGroup.querySelectorAll(".sample-chip"));
  if (!visibleChips.length) return;
  const currentActive = visibleChips.find((c) => c.classList.contains("active"));
  const pool = visibleChips.filter((c) => c !== currentActive);
  const target = (pool.length ? pool : visibleChips)[Math.floor(Math.random() * (pool.length || visibleChips.length))];
  selectSample(target);
});

const thresholdInput = document.getElementById("span-threshold");
const thresholdVal = document.getElementById("span-threshold-val");
if (thresholdInput) {
  thresholdInput.addEventListener("input", () => {
    spanThreshold = Number(thresholdInput.value);
    if (thresholdVal) thresholdVal.textContent = spanThreshold.toFixed(2);
  });
}

// Fetch model READMEs (metrics + training details) for the "Models & tasks" panel.
async function loadReadme(repo, elId) {
  const el = document.getElementById(elId);
  if (!el) return;
  try {
    const resp = await fetch(`https://huggingface.co/${repo}/resolve/main/README.md`);
    if (!resp.ok) throw new Error(String(resp.status));
    const md = await resp.text();
    const metrics = md
      .split("\n")
      .filter((l) => l.includes("|") && (/\d/.test(l) || l.startsWith("|")))
      .join("\n");
    el.textContent = metrics || md.slice(0, 1500);
  } catch {
    el.textContent = "README unavailable offline.";
  }
}
loadReadme("rafmacalaba/gliner_datause_extended", "meta-gliner");
loadReadme("rafmacalaba/lfm2.5-350M-datause-multitask", "meta-multitask");