import { prebuiltAppConfig, CreateMLCEngine } from "@charliefruan/web-llm";
import hljs from "highlight.js";
import ace from "ace-builds";

// required for ace to resolve module correctly
require("ace-builds/src-noconflict/mode-javascript");
require("ace-builds/webpack-resolver");

// DO NOT REMOVE
// Required for user input type definition to be eval
const { Type } = require('@sinclair/typebox');

let engine = null;

const availableModels = prebuiltAppConfig.model_list
  .filter(
    (m) =>
      m.model_id.startsWith("Llama-3-8B") ||
      m.model_id.startsWith("Llama-3-70B") ||
      m.model_id.startsWith("Llama-3.1-8B") ||
      m.model_id.startsWith("Llama-3.1-70B") ||
      m.model_id.startsWith("Llama-3.2-3B") ||
      m.model_id.startsWith("Hermes-2") ||
      m.model_id.startsWith("Hermes-3") ||
      m.model_id.startsWith("Phi-3")
  )
  .map((m) => m.model_id);
let selectedModel = availableModels[0];

availableModels.forEach((modelId) => {
  const option = document.createElement("option");
  option.value = modelId;
  option.textContent = modelId;
  document.getElementById("model-selection").appendChild(option);
});
document.getElementById("model-selection").value = selectedModel;
document.getElementById("model-selection").onchange = (e) => {
  selectedModel = e.target.value;
  engine = null;
};
document.getElementById(
  "prompt"
).value = `Hermione Granger is a character in Harry Potter. Please fill in the following information about this character in JSON format.
Name is a string of character name.
House is one of Gryffindor, Hufflepuff, Ravenclaw, Slytherin.
Blood status is one of Pure-blood, Half-blood, Muggle-born.
Occupation is one of Student, Professor, Ministry of Magic, Other.
Wand is an object with wood, core, and length.
Alive is a boolean.
Patronus is a string.
`;

// JSON editor setup
const editor = ace.edit("schema", {
  //   mode: "ace/mode/javascript",
  mode: "ace/mode/javascript",
  theme: "ace/theme/github",
  wrap: true,
});
editor.setTheme("ace/theme/github");
editor.setValue(`Type.Object({
  "name": Type.String(),
  "house": Type.Enum({
    "Gryffindor": "Gryffindor",
    "Hufflepuff": "Hufflepuff",
    "Ravenclaw": "Ravenclaw",
    "Slytherin": "Slytherin",
  }),
  "blood_status": Type.Enum({
    "Pure-blood": "Pure-blood",
    "Half-blood": "Half-blood",
    "Muggle-born": "Muggle-born",
  }),
  "occupation": Type.Enum({
    "Student": "Student",
    "Professor": "Professor",
    "Ministry of Magic": "Ministry of Magic",
    "Other": "Other",
  }),
  "wand": Type.Object({
    "wood": Type.String(),
    "core": Type.String(),
    "length": Type.Number(),
  }),
  "alive": Type.Boolean(),
  "patronus": Type.String(),
})`);

// Generate button
document.getElementById("generate").onclick = async () => {
  const schemaInput = editor.getValue();
  let T;
  try {
    T = eval(schemaInput);
  } catch (e) {
    console.error("Invalid schema", e);
    return;
  }
  const schema = JSON.stringify(T);

  if (!engine) {
    engine = await CreateMLCEngine(selectedModel, {
      initProgressCallback: (progress) => {
        console.log(progress);
        document.getElementById("output").textContent = progress.text;
      },
    });
  }
  const request = {
    stream: true,
    stream_options: { include_usage: true },
    messages: [
      {
        role: "user",
        content: document.getElementById("prompt").value,
      },
    ],
    max_tokens: 128,
    response_format: {
      type: "json_object",
      schema: schema,
    },
  };

  let curMessage = "";
  let usage = null;
  const generator = await engine.chatCompletion(request);
  for await (const chunk of generator) {
    const curDelta = chunk.choices[0]?.delta.content;
    if (curDelta) {
      curMessage += curDelta;
    }
    if (chunk.usage) {
      usage = chunk.usage;
    }
    document.getElementById("output").textContent = curMessage;
  }
  const finalMessage = await engine.getMessage();
  console.log(finalMessage);
  if (hljs) {
    document.getElementById("output").innerHTML = hljs.highlight(finalMessage, {
      language: "json",
    }).value;
  } else {
    document.getElementById("output").textContent = finalMessage;
  }
  if (usage) {
    const stats = usage['extra'];
    document.getElementById("stats").textContent = `Prefill: ${stats['prefill_tokens_per_s'].toFixed(2)}, Decode: ${stats['decode_tokens_per_s'].toFixed(2)}, Grammar Init: ${stats['grammar_init_ms'].toFixed(2)}, Grammar Overhead Per Token: ${stats['grammar_per_token_ms'].toFixed(2)}`;
    document.getElementById("stats").classList.remove("hidden");
  }
};