import { Tokenizer } from "https://nobsdelivr.private.coffee/npm/@huggingface/tokenizers@0.1.3";

if (localStorage.getItem("tokenizerCache")) {
  var tokenizerCache = JSON.parse(localStorage.tokenizerCache);
} else {
  var tokenizerCache = {};
  localStorage.tokenizerCache = "{}";
}
var tokenizerJSON;

let tokenizeButton = document.querySelector("#tokenize");
tokenizeButton.addEventListener("click", async () => {
  tokenizeButton.disabled = true;
  let modelId = document.querySelector("#model").value;
  
  if (modelId in tokenizerCache) {
    tokenizerJSON = tokenizerCache[modelId];
  }
  else {
    let tokenizerJSONSocket = await fetch(`https://huggingface.co/${modelId}/resolve/main/tokenizer.json`);
    tokenizerJSON = await tokenizerJSONSocket.json();
    tokenizerCache[modelId] = tokenizerJSON;
    localStorage.tokenizerCache = JSON.stringify(tokenizerCache);
  }
  
  let tokenizer = new Tokenizer(tokenizerJSON, {});
  
  let text = document.querySelector("#text").value;
  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);
  });
  
  tokenizeButton.disabled = false;
});

document.querySelector("#clear-cache").addEventListener("click", () => {
  tokenizerCache = {};
  localStorage.tokenizerCache = "{}";
});