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

// Required for ace to resolve the 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;
let useCustomGrammar = false;

document.addEventListener("DOMContentLoaded", () => {
  // Ensure elements are loaded before using them
  const grammarSelection = document.getElementById("grammar-selection");
  const ebnfContainer = document.getElementById("ebnf-grammar-container");
  const schemaContainer = document.getElementById("schema-container");
  const modelSelection = document.getElementById("model-selection");
  const ebnfTextarea = document.getElementById("ebnf-grammar");
  const promptTextarea = document.getElementById("prompt");
  const outputDiv = document.getElementById("output");
  const statsParagraph = document.getElementById("stats");

  // Initialize the custom grammar textarea
  ebnfTextarea.value = String.raw`main ::= basic_array | basic_object
basic_any ::= basic_number | basic_string | basic_boolean | basic_null | basic_array | basic_object
basic_integer ::= ("0" | "-"? [1-9] [0-9]*) ".0"?
basic_number ::= ("0" | "-"? [1-9] [0-9]*) ("." [0-9]+)? ([eE] [+-]? [0-9]+)?
basic_string ::= (([\"] basic_string_1 [\"]))
basic_string_1 ::= "" | [^"\\\x00-\x1F] basic_string_1 | "\\" escape basic_string_1
escape ::= ["\\/bfnrt] | "u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9]
basic_boolean ::= "true" | "false"
basic_null ::= "null"
basic_array ::= "[" ("" | ws basic_any (ws "," ws basic_any)*) ws "]"
basic_object ::= "{" ("" | ws basic_string ws ":" ws basic_any ( ws "," ws basic_string ws ":" ws basic_any)*) ws "}"
ws ::= [\n\t]*`;

  // Handle grammar selection changes
  grammarSelection.onchange = (ev) => {
    console.log("Grammar selection changed:", ev.target.value);
    if (ev.target.value === "json") {
      ebnfContainer.classList.add("hidden");
      schemaContainer.classList.remove("hidden");
      useCustomGrammar = false;
    } else {
      ebnfContainer.classList.remove("hidden");
      schemaContainer.classList.add("hidden");
      useCustomGrammar = true;
    }
  };

  // Populate model selection dropdown
  const availableModels = prebuiltAppConfig.model_list
    .filter(
      (m) =>
        m.model_id.startsWith("Llama-3") ||
        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;
    modelSelection.appendChild(option);
  });

  modelSelection.value = selectedModel;

  modelSelection.onchange = (e) => {
    selectedModel = e.target.value;
    engine = null; // Reset the engine when the model changes
  };

  // JSON editor setup with Ace
  const editor = ace.edit("schema", {
    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(),
  })`);

  // Set initial prompt
  promptTextarea.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.
`;
  // Generate button click handler
  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);
          outputDiv.textContent = progress.text;
        },
      });
    }

    const request = {
      stream: true,
      stream_options: { include_usage: true },
      messages: [{ role: "user", content: promptTextarea.value }],
      max_tokens: 128,
      response_format: useCustomGrammar
        ? { type: "grammar", grammar: ebnfTextarea.value }
        : { 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;
      outputDiv.textContent = curMessage;
    }

    const finalMessage = await engine.getMessage();
    outputDiv.innerHTML = hljs.highlight(finalMessage, {
      language: "json",
    }).value;

    if (usage) {
      const statsTextParts = [];
      if (usage.extra.prefill_tokens_per_s) {
        statsTextParts.push(`Prefill: ${usage.extra.prefill_tokens_per_s.toFixed(
          1
        )} tok/s`);
      }
      if (usage.extra.decode_tokens_per_s) {
        statsTextParts.push(`Decode: ${usage.extra.decode_tokens_per_s.toFixed(
          1
        )} tok/s`);
      }
      if (usage.extra.grammar_init_ms) {
        statsTextParts.push(`Grammar Init Overhead: ${usage.extra.grammar_init_ms.toFixed(
          0
        )} ms`);
      }
      if (usage.extra.grammar_per_token_ms) {
        statsTextParts.push(`Grammar Per Token Overhead: ${usage.extra.grammar_per_token_ms.toFixed(
          2
        )} ms`);
      }
      statsParagraph.textContent = statsTextParts.join(", ");
      statsParagraph.classList.remove("hidden");
    }
  };
});