import { Tokenizer } from "./tokenizers.mjs";

// IndexedDB helper for caching tokenizer JSON.
// IndexedDB can store large blobs, unlike localStorage's ~5MB cap.
const dbName = "tokenizerCacheDB";
const storeName = "tokenizerCache";

function openDB() {
  return new Promise((resolve, reject) => {
    const request = indexedDB.open(dbName, 1);

    request.onupgradeneeded = () => {
      request.result.createObjectStore(storeName);
    };

    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
  });
}

let dbPromise;

async function getCache(modelId) {
  if (!dbPromise) dbPromise = openDB();
  const db = await dbPromise;

  return new Promise((resolve, reject) => {
    const tx = db.transaction(storeName, "readonly");
    const get = tx.objectStore(storeName).get(modelId);

    get.onsuccess = () => resolve(get.result);
    get.onerror = () => reject(get.error);
  });
}

async function setCache(modelId, tokenizerJSON) {
  if (!dbPromise) dbPromise = openDB();
  const db = await dbPromise;

  return new Promise((resolve, reject) => {
    const tx = db.transaction(storeName, "readwrite");
    tx.objectStore(storeName).put(tokenizerJSON, modelId);

    tx.oncomplete = () => resolve();
    tx.onerror = () => reject(tx.error);
  });
}

async function clearCache() {
  if (!dbPromise) dbPromise = openDB();
  const db = await dbPromise;

  return new Promise((resolve, reject) => {
    const tx = db.transaction(storeName, "readwrite");
    tx.objectStore(storeName).clear();

    tx.oncomplete = () => resolve();
    tx.onerror = () => reject(tx.error);
  });
}

// Fetch a tokenizer.json while reporting download progress via the
// <progress> element. Uses the response stream so we can count bytes.
async function fetchTokenizer(modelId) {
  const response = await fetch(`https://huggingface.co/${modelId}/resolve/main/tokenizer.json`);
  if (!response.ok) {
    throw new Error(`Download failed (HTTP ${response.status})`);
  }

  const total = Number(response.headers.get("content-length")) || 0;
  const reader = response.body.getReader();
  const chunks = [];
  let received = 0;

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    received += value.length;
    if (total) showProgress((received / total) * 100);
  }

  const bytes = new Uint8Array(received);
  let offset = 0;
  for (const chunk of chunks) {
    bytes.set(chunk, offset);
    offset += chunk.length;
  }

  return JSON.parse(new TextDecoder().decode(bytes));
}

const statusEl = document.querySelector("#status");
const progressEl = document.querySelector("#progress");
const progressWrapEl = document.querySelector("#progress-wrap");

function setStatus(text, isError = false) {
  statusEl.textContent = text;
  statusEl.hidden = !text;
  statusEl.classList.toggle("error", isError);
}

function showProgress(percent) {
  progressEl.value = percent;
  progressWrapEl.hidden = false;
}

function hideProgress() {
  progressEl.value = 0;
  progressWrapEl.hidden = true;
}

let tokenizeButton = document.querySelector("#tokenize");
tokenizeButton.addEventListener("click", async () => {
  tokenizeButton.disabled = true;
  setStatus("");
  hideProgress();

  const modelId = document.querySelector("#model").value;

  try {
    const text = document.querySelector("#text").value;

    let tokenizerJSON;
    setStatus(`Loading tokenizer for ${modelId}…`);
    try {
      tokenizerJSON = await getCache(modelId);
    } catch (err) {
      console.error("Failed to read from IndexedDB:", err);
    }

    if (tokenizerJSON === undefined) {
      setStatus(`Downloading ${modelId} tokenizer…`);
      tokenizerJSON = await fetchTokenizer(modelId);
      try {
        await setCache(modelId, tokenizerJSON);
      } catch (err) {
        console.error("Failed to write to IndexedDB:", err);
      }
    } else {
      setStatus(`Loading cached tokenizer…`);
    }

    hideProgress();
    setStatus("Building tokenizer…");
    let tokenizer = new Tokenizer(tokenizerJSON, {});

    setStatus("Tokenizing…");
    // Yield to the event loop so the browser can paint the status message
    // before the synchronous encode/decode blocks the main thread. (A rAF
    // alone runs *before* paint, so it wouldn't show.)
    await new Promise((resolve) => setTimeout(resolve, 50));
    let tokenized = tokenizer.encode(text);

    document.querySelector("#totalLength").textContent = `Total length: ${tokenized.ids.length}`;

    let output = document.querySelector("#output");
    output.contentDocument.body.innerHTML = "";

    let colors = ["#ff00007f", "#00ff007f", "#ffff007f", "#0000ff7f"].sort(() => Math.random() - 0.5);

    tokenized.ids.forEach((token, index) => {
      let content = document.createElement("span");

      content.textContent = tokenizer.decode([token]);
      content.style.backgroundColor = colors[index % colors.length];
      content.style.fontFamily = "monospace";
      content.style.whiteSpace = "pre-wrap";

      content.title = token.toString();

      output.contentDocument.body.appendChild(content);
    });

    // Clear feedback once done (no lingering completion message).
    setStatus("");
  } catch (err) {
    console.error(err);
    hideProgress();
    setStatus(`Error: ${err.message}. Try clearing the cache.`, true);
  } finally {
    tokenizeButton.disabled = false;
  }
});

document.querySelector("#clear-cache").addEventListener("click", async () => {
  try {
    await clearCache();
    setStatus("Cache cleared.");
  } catch (err) {
    console.error("Failed to clear cache:", err);
    setStatus(`Error clearing cache: ${err.message}`, true);
  }
});