import * as webllm from "https://esm.run/@mlc-ai/web-llm";
import hljs from "https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/es/highlight.min.js";
import { Type } from "./lib/typebox/index.mjs";

let engine = null;

const availableModels = webllm.prebuiltAppConfig.model_list
  .filter(
    (m) =>
      m.model_id.startsWith("Llama-3") ||
      m.model_id.startsWith("Hermes-2") ||
      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 webllm.CreateMLCEngine(selectedModel, {
      initProgressCallback: (progress) => {
        console.log(progress);
        document.getElementById("output").textContent = progress.text;
      },
    });
  }
  const request = {
    stream: true,
    messages: [
      {
        role: "user",
        content: document.getElementById("prompt").value,
      },
    ],
    max_tokens: 128,
    response_format: {
      type: "json_object",
      schema: schema,
    },
  };

  let curMessage = "";
  const generator = await engine.chatCompletion(request);
  for await (const chunk of generator) {
    const curDelta = chunk.choices[0]?.delta.content;
    if (curDelta) {
      curMessage += curDelta;
    }
    console.log(curMessage);
    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;
  }
};