import { TOKENIZER_OPTIONS } from "./tokenizers.js";

// ─── Token color palette ──────────────────────────────────────────────────────
// Reference CSS custom properties so dark-mode overrides apply automatically.
// The index wraps modulo 12; values are defined in style.css as --tok-N.

const NUM_TOKEN_COLORS = 12;

function tokenColorVar(index) {
  return `var(--tok-${index % NUM_TOKEN_COLORS})`;
}

const DEBOUNCE_MS = 350;

// ─── State ────────────────────────────────────────────────────────────────────

/** @type {Map<number, PanelState>} */
const panels = new Map();
let panelCounter = 0;

/**
 * @typedef {{
 *   modelId: string,
 *   name: string,
 *   worker: Worker,
 *   panel: HTMLElement,
 *   textMirror: HTMLElement,
 *   statsEl: HTMLElement,
 *   tbody: HTMLElement,
 *   tokenCount: number,
 * }} PanelState
 */

// ─── DOM refs ─────────────────────────────────────────────────────────────────

const textInput = /** @type {HTMLTextAreaElement} */ (
  document.getElementById("text-input")
);
const selectEl = /** @type {HTMLSelectElement}   */ (
  document.getElementById("tokenizer-select")
);
const addBtn = /** @type {HTMLButtonElement}   */ (
  document.getElementById("add-btn")
);
const panelsEl = /** @type {HTMLElement}         */ (
  document.getElementById("panels")
);
const emptyState = /** @type {HTMLElement}         */ (
  document.getElementById("empty-state")
);
const highlightStyleEl = /** @type {HTMLStyleElement}    */ (
  document.getElementById("highlight-styles")
);

// ─── Populate dropdown ────────────────────────────────────────────────────────

for (const [modelId, name] of Object.entries(TOKENIZER_OPTIONS)) {
  if (!modelId) continue;
  const opt = document.createElement("option");
  opt.value = modelId;
  opt.textContent = name;
  selectEl.appendChild(opt);
}

selectEl.addEventListener("change", () => {
  addBtn.disabled = !selectEl.value;
});

addBtn.addEventListener("click", () => {
  const modelId = selectEl.value;
  if (!modelId) return;
  const name = TOKENIZER_OPTIONS[modelId] ?? modelId;
  addPanel(modelId, name);
  selectEl.value = "";
  addBtn.disabled = true;
});

// ─── Text input ───────────────────────────────────────────────────────────────

let debounceTimer = null;

textInput.addEventListener("input", () => {
  clearTimeout(debounceTimer);
  debounceTimer = setTimeout(retokenizeAll, DEBOUNCE_MS);
});

function retokenizeAll() {
  const text = textInput.value;
  for (const id of panels.keys()) {
    runTokenizer(id, text);
  }
}

// ─── Panel lifecycle ──────────────────────────────────────────────────────────

function addPanel(modelId, name) {
  const id = panelCounter++;

  const panel = document.createElement("div");
  panel.className = "panel loading";
  panel.setAttribute("role", "listitem");
  panel.dataset.panelId = String(id);
  panel.innerHTML = panelTemplate(name, modelId);

  panel
    .querySelector(".remove-btn")
    .addEventListener("click", () => removePanel(id));

  panelsEl.appendChild(panel);
  updateEmptyState();

  const worker = new Worker("js/worker.js", { type: "module" });
  worker.addEventListener("message", (ev) => onWorkerMessage(id, ev.data));
  worker.addEventListener("error", (ev) =>
    onWorkerError(id, ev.message ?? "Worker error"),
  );

  panels.set(id, {
    modelId,
    name,
    worker,
    panel,
    textMirror: panel.querySelector(".panel-text"),
    statsEl: panel.querySelector(".panel-stats"),
    tbody: panel.querySelector("tbody"),
    tokenCount: 0,
  });

  runTokenizer(id, textInput.value);
}

function removePanel(id) {
  const state = panels.get(id);
  if (!state) return;
  clearPanelHighlights(id, state.tokenCount);
  state.worker.terminate();
  state.panel.remove();
  panels.delete(id);
  updateEmptyState();
}

function updateEmptyState() {
  emptyState.classList.toggle("hidden", panels.size > 0);
}

function panelTemplate(name, modelId) {
  return `
    <div class="panel-header">
      <a class="panel-title" href="https://huggingface.co/${esc(modelId)}" target="_blank" rel="noopener noreferrer" title="${esc(modelId)}">${esc(name)}</a>
      <button class="remove-btn" type="button" aria-label="Remove ${esc(name)}">Remove</button>
    </div>
    <div class="panel-loading">
      <span class="spinner" aria-hidden="true"></span>
      Loading tokenizer…
    </div>
    <div class="panel-text" aria-label="Tokenized text"></div>
    <div class="panel-stats"></div>
    <div class="panel-table-wrap">
      <table aria-label="Token list">
        <thead>
          <tr>
            <th></th>
            <th>#</th>
            <th>Token</th>
            <th>ID</th>
          </tr>
        </thead>
        <tbody></tbody>
      </table>
    </div>`;
}

// ─── Worker communication ─────────────────────────────────────────────────────

function runTokenizer(id, text) {
  const state = panels.get(id);
  if (!state) return;
  state.panel.classList.add("loading");
  state.worker.postMessage({ model_id: state.modelId, text });
}

/**
 * @param {number} id
 * @param {{ token_ids: number[], decoded: string[], margins: number[] }} data
 */
function onWorkerMessage(id, data) {
  const state = panels.get(id);
  if (!state) return;

  state.panel.classList.remove("loading");

  const { token_ids, decoded, margins } = data;
  const text = textInput.value;

  // Clear stale highlights
  clearPanelHighlights(id, state.tokenCount);
  state.tokenCount = token_ids.length;

  // Sync mirror text (single text node — required for Range offsets)
  state.textMirror.textContent = text;

  // Compute per-token character offsets in the original string
  const offsets = computeOffsets(text, decoded, margins ?? []);

  // Paint highlights via CSS Custom Highlight API
  paintHighlights(id, state.textMirror, offsets);

  // Stats bar
  const ratio =
    text.length > 0 ? (text.length / token_ids.length).toFixed(2) : "—";
  state.statsEl.innerHTML =
    `<span>Tokens: <strong>${token_ids.length}</strong></span>` +
    `<span>Chars/token: <strong>${ratio}</strong></span>` +
    `<span>Characters: <strong>${text.length}</strong></span>`;

  // Table
  renderTable(id, state, token_ids, decoded, offsets);
}

function onWorkerError(id, msg) {
  const state = panels.get(id);
  if (!state) return;
  state.panel.classList.remove("loading");
  state.textMirror.textContent = `! ${msg}`;
}

// ─── Offset computation ───────────────────────────────────────────────────────

/**
 * Walk `decoded` tokens greedily left-to-right against `text`, returning
 * [{start, end}|null] for each token.
 *
 * The HuggingFace worker decodes each token individually.  Metaspace converts
 * the leading ▁ (U+2581) back to a space, so `decoded` values should
 * concatenate to the original text (modulo special tokens like <s>, </s>).
 *
 * @param {string} text
 * @param {string[]} decoded
 * @param {number[]} margins  — BERT: margin>0 means a word boundary space precedes
 * @returns {Array<{start:number,end:number}|null>}
 */
function computeOffsets(text, decoded, margins) {
  const offsets = [];
  let cursor = 0;

  for (let i = 0; i < decoded.length; i++) {
    const tok = decoded[i] ?? "";

    if (tok.length === 0) {
      offsets.push(null);
      continue;
    }

    // BERT word-boundary space: margin > 0 means we should skip a space
    if (margins[i] > 0 && text[cursor] === " ") {
      cursor++;
    }

    // Greedy forward search from cursor
    const idx = text.indexOf(tok, cursor);
    if (idx === -1) {
      // Special token or mismatch — no visible range
      offsets.push(null);
      continue;
    }

    offsets.push({ start: idx, end: idx + tok.length });
    cursor = idx + tok.length;
  }

  return offsets;
}

// ─── CSS Custom Highlight API ─────────────────────────────────────────────────

// Track which ::highlight() rules have been injected to avoid duplication
const injectedRules = new Set();

/**
 * Ensure ::highlight(name) rule exists in the shared <style> element.
 * @param {string} name
 * @param {string} bgColor
 */
function ensureRule(name, bgColor) {
  if (injectedRules.has(name)) return;
  injectedRules.add(name);
  const sheet = highlightStyleEl.sheet;
  sheet.insertRule(
    `::highlight(${name}) { background-color: ${bgColor}; color: var(--tok-text); }`,
    sheet.cssRules.length,
  );
}

/**
 * Ensure the hover ::highlight() rule exists for a panel.
 * @param {number} panelId
 */
function ensureHoverRule(panelId) {
  const name = `p${panelId}-hover`;
  if (injectedRules.has(name)) return;
  injectedRules.add(name);
  const sheet = highlightStyleEl.sheet;
  sheet.insertRule(
    `::highlight(${name}) { outline: 2px solid var(--tok-hover); border-radius: 2px; }`,
    sheet.cssRules.length,
  );
}

/**
 * Register one CSS Highlight per token offset.
 * Name: `p{panelId}-t{tokenIndex}`
 * Color: var(--tok-N) from CSS
 */
function paintHighlights(panelId, mirrorEl, offsets) {
  if (!CSS.highlights) return;

  const textNode = mirrorEl.firstChild;
  if (!textNode) return;

  for (let i = 0; i < offsets.length; i++) {
    const off = offsets[i];
    if (!off) continue;

    const name = `p${panelId}-t${i}`;

    ensureRule(name, tokenColorVar(i));

    const range = new Range();
    range.setStart(textNode, off.start);
    range.setEnd(textNode, off.end);
    CSS.highlights.set(name, new Highlight(range));
  }
}

/**
 * Remove all named highlights for a panel.
 */
function clearPanelHighlights(panelId, tokenCount) {
  if (!CSS.highlights) return;
  for (let i = 0; i < tokenCount; i++) {
    CSS.highlights.delete(`p${panelId}-t${i}`);
  }
  CSS.highlights.delete(`p${panelId}-hover`);
}

// ─── Token table ──────────────────────────────────────────────────────────────

/**
 * @param {number} panelId
 * @param {PanelState} state
 * @param {number[]} token_ids
 * @param {string[]} decoded
 * @param {Array<{start:number,end:number}|null>} offsets
 */
function renderTable(panelId, state, token_ids, decoded, offsets) {
  const { tbody, textMirror } = state;
  tbody.innerHTML = "";

  ensureHoverRule(panelId);

  const fragment = document.createDocumentFragment();

  for (let i = 0; i < token_ids.length; i++) {
    const color = tokenColorVar(i);
    const off = offsets[i];

    const tr = document.createElement("tr");
    tr.style.setProperty("--row-color", color);
    tr.dataset.tokenIndex = String(i);

    // Color swatch
    const tdSwatch = document.createElement("td");
    tdSwatch.className = "color-swatch";
    const swatch = document.createElement("span");
    swatch.style.background = color;
    tdSwatch.appendChild(swatch);

    // Index
    const tdIdx = document.createElement("td");
    tdIdx.textContent = String(i + 1);

    // Token text
    const tdTok = document.createElement("td");
    tdTok.className = "tok-cell";
    tdTok.textContent = decoded[i] ?? "";

    // Token ID
    const tdId = document.createElement("td");
    tdId.className = "id-cell";
    tdId.textContent = String(token_ids[i]);

    tr.append(tdSwatch, tdIdx, tdTok, tdId);

    // Hover: show distinct outline highlight on the mirrored text
    if (off) {
      tr.addEventListener("mouseenter", () => {
        tr.classList.add("active");
        if (!CSS.highlights) return;
        const node = textMirror.firstChild;
        if (!node) return;
        const r = new Range();
        r.setStart(node, off.start);
        r.setEnd(node, off.end);
        CSS.highlights.set(`p${panelId}-hover`, new Highlight(r));
      });
      tr.addEventListener("mouseleave", () => {
        tr.classList.remove("active");
        CSS.highlights?.delete(`p${panelId}-hover`);
      });
    }

    fragment.appendChild(tr);
  }

  tbody.appendChild(fragment);
}

// ─── Utility ──────────────────────────────────────────────────────────────────

function esc(str) {
  return str
    .replace(/&/g, "&amp;")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/"/g, "&quot;");
}