import { pipeline, env } from "@xenova/transformers";

const MODEL_ID = "onnx-community/orpheus-3b-0.1-ft-ONNX";

const modal       = document.getElementById("modal");
const modalText   = document.getElementById("modal-text");
const progressBar = document.getElementById("progress-fill");
const inputEl     = document.getElementById("input-text");
const btnEl       = document.getElementById("generate-btn");
const loadingEl   = document.getElementById("loading-text");
const outputEl    = document.getElementById("output-text");

let generator = null;

// Set up ONNX WASM paths (optional)
env.backends.onnx.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/@xenova/transformers@2.17.2/dist/";
env.useWasmCPU();

btnEl.addEventListener("click", async () => {
  const prompt = inputEl.value.trim();
  if (!prompt) return;

  btnEl.disabled = true;
  outputEl.textContent = "";
  loadingEl.textContent = "";

  // Show downloading modal
  modal.classList.remove("hidden");
  progressBar.style.width = "0%";
  modalText.textContent = "Loading model... 0%";

  // Load (or reuse) pipeline with progress callback :contentReference[oaicite:0]{index=0}
  if (!generator) {
    generator = await pipeline("text2text-generation", MODEL_ID, {
      progress_callback: (args) => {
        if (args.status === "progress") {
          const pct = Math.floor(args.progress * 100);
          progressBar.style.width = pct + "%";
          modalText.textContent = `Loading model... ${pct}%`;
        }
      },
    });
  }

  // Hide modal
  modalText.textContent = "Model ready!";
  await new Promise((r) => setTimeout(r, 300));
  modal.classList.add("hidden");

  // Stream tokens as they generate :contentReference[oaicite:1]{index=1}
  loadingEl.classList.remove("hidden");
  await generator(prompt, {
    max_length: 256,
    callback_function: (beams) => {
      // beams may be an array of { token, text, ... }
      // append the latest token or text
      const tok = typeof beams === "string"
        ? beams
        : (beams[0]?.text || beams[0]?.token || "");
      loadingEl.textContent += tok;
    },
  });

  // Final output from the pipeline
  const [res] = await generator(prompt, { max_length: 0 });
  loadingEl.classList.add("hidden");
  outputEl.textContent = res.generated_text;

  btnEl.disabled