import { env, pipeline } from "@huggingface/transformers";
import "./style.css";

const MODEL_ID = "Reza2kn/MiniCPM5-1B-ONNX-Web";

const output = document.querySelector("#output");
const run = document.querySelector("#run");
const prompt = document.querySelector("#prompt");
const maxTokens = document.querySelector("#maxTokens");
const temperature = document.querySelector("#temperature");
const backend = document.querySelector("#backend");

env.allowLocalModels = false;
env.backends.onnx.wasm.numThreads = Math.min(4, navigator.hardwareConcurrency || 4);

let generatorPromise = null;

function setStatus(text) {
  output.textContent = text;
}

async function getGenerator() {
  if (!generatorPromise) {
    setStatus(`Loading ${MODEL_ID} with ${backend.value}...`);
    generatorPromise = pipeline("text-generation", MODEL_ID, {
      device: backend.value,
      dtype: "q4",
    });
  }
  return generatorPromise;
}

backend.addEventListener("change", () => {
  generatorPromise = null;
  setStatus("Backend changed. Run again to reload.");
});

run.addEventListener("click", async () => {
  run.disabled = true;
  try {
    const generator = await getGenerator();
    setStatus("Generating...");
    const result = await generator(prompt.value, {
      max_new_tokens: Number(maxTokens.value),
      temperature: Number(temperature.value),
      do_sample: Number(temperature.value) > 0,
    });
    setStatus(result?.[0]?.generated_text ?? JSON.stringify(result, null, 2));
  } catch (error) {
    setStatus(error.stack || error.message || String(error));
  } finally {
    run.disabled = false;
  }
});

setStatus("Ready. Upload the ONNX-Web artifact, then run.");