import * as ort from "onnxruntime-web";
import "./style.css";

const TASK_LABELS = ["general_chat", "writing", "translation", "summarization", "research", "coding", "mathematics", "document_analysis", "high_stakes"];
const COMPLEXITY_LABELS = ["low", "medium", "high"];
const WORD_BINS = 1024;
const CHAR_BINS = 512;
const INPUT_SIZE = 1544;

function fnv1a(text) {
  let hash = 2166136261;
  for (let i = 0; i < text.length; i += 1) {
    hash ^= text.charCodeAt(i);
    hash = Math.imul(hash, 16777619);
  }
  return hash >>> 0;
}

function buildFeatures(text) {
  const source = String(text || "");
  const lower = source.toLocaleLowerCase();
  const values = new Float32Array(INPUT_SIZE);
  const words = lower.match(/[\p{L}\p{N}_]+/gu) || [];
  const grams = [...words];
  for (let i = 0; i + 1 < words.length; i += 1) grams.push(`${words[i]}_${words[i + 1]}`);
  for (const gram of grams) values[fnv1a(gram) % WORD_BINS] += 1;

  const compact = lower.replace(/\s+/g, " ");
  for (let i = 0; i + 2 < compact.length; i += 1) {
    values[WORD_BINS + (fnv1a(compact.slice(i, i + 3)) % CHAR_BINS)] += 0.25;
  }

  const sparseCount = WORD_BINS + CHAR_BINS;
  let norm = 0;
  for (let i = 0; i < sparseCount; i += 1) norm += values[i] * values[i];
  norm = Math.sqrt(norm);
  if (norm > 0) for (let i = 0; i < sparseCount; i += 1) values[i] /= norm;

  const base = sparseCount;
  values[base] = Math.min(source.length, 4000) / 4000;
  values[base + 1] = Math.min(words.length, 800) / 800;
  values[base + 2] = Math.min((source.match(/\?/g) || []).length, 10) / 10;
  values[base + 3] = Math.min((source.match(/\n/g) || []).length, 30) / 30;
  values[base + 4] = Math.min((source.match(/[{}[\]();=<>/]/g) || []).length, 100) / 100;
  values[base + 5] = /https?:\/\/|www\./i.test(source) ? 1 : 0;
  values[base + 6] = /\b(pdf|document|dokument|file|súbor|attachment|príloha)\b/i.test(source) ? 1 : 0;
  values[base + 7] = 1;
  return values;
}

function argMax(values) {
  let index = 0;
  for (let i = 1; i < values.length; i += 1) if (values[i] > values[index]) index = i;
  return index;
}

const result = document.querySelector("#result");
const modelUrl = new URL("../../../model.onnx", import.meta.url).href;
ort.env.wasm.numThreads = 1;
const session = await ort.InferenceSession.create(modelUrl, { executionProviders: ["wasm"] });
result.textContent = "Model ready.";

document.querySelector("#classify").addEventListener("click", async () => {
  const text = document.querySelector("#prompt").value;
  const features = buildFeatures(text);
  const outputs = await session.run({ features: new ort.Tensor("float32", features, [1, INPUT_SIZE]) });
  const task = Array.from(outputs.task_probabilities.data);
  const complexity = Array.from(outputs.complexity_probabilities.data);
  const taskIndex = argMax(task);
  const complexityIndex = argMax(complexity);
  result.textContent = JSON.stringify({
    task: { label: TASK_LABELS[taskIndex], confidence: task[taskIndex], probabilities: Object.fromEntries(TASK_LABELS.map((label, i) => [label, task[i]])) },
    complexity: { label: COMPLEXITY_LABELS[complexityIndex], confidence: complexity[complexityIndex], probabilities: Object.fromEntries(COMPLEXITY_LABELS.map((label, i) => [label, complexity[i]])) }
  }, null, 2);
});