import type { HighlighterCore, LanguageInput, ThemedToken } from "shiki/core";

type LanguageLoader = () => Promise<{ default: LanguageInput }>;

const languageLoaders: Record<string, LanguageLoader> = {
  bash: () => import("shiki/langs/bash.mjs"),
  c: () => import("shiki/langs/c.mjs"),
  cpp: () => import("shiki/langs/cpp.mjs"),
  csharp: () => import("shiki/langs/csharp.mjs"),
  css: () => import("shiki/langs/css.mjs"),
  dockerfile: () => import("shiki/langs/dockerfile.mjs"),
  go: () => import("shiki/langs/go.mjs"),
  html: () => import("shiki/langs/html.mjs"),
  java: () => import("shiki/langs/java.mjs"),
  javascript: () => import("shiki/langs/javascript.mjs"),
  json: () => import("shiki/langs/json.mjs"),
  jsx: () => import("shiki/langs/jsx.mjs"),
  kotlin: () => import("shiki/langs/kotlin.mjs"),
  markdown: () => import("shiki/langs/markdown.mjs"),
  php: () => import("shiki/langs/php.mjs"),
  powershell: () => import("shiki/langs/powershell.mjs"),
  python: () => import("shiki/langs/python.mjs"),
  ruby: () => import("shiki/langs/ruby.mjs"),
  rust: () => import("shiki/langs/rust.mjs"),
  scss: () => import("shiki/langs/scss.mjs"),
  sql: () => import("shiki/langs/sql.mjs"),
  swift: () => import("shiki/langs/swift.mjs"),
  toml: () => import("shiki/langs/toml.mjs"),
  tsx: () => import("shiki/langs/tsx.mjs"),
  typescript: () => import("shiki/langs/typescript.mjs"),
  vue: () => import("shiki/langs/vue.mjs"),
  xml: () => import("shiki/langs/xml.mjs"),
  yaml: () => import("shiki/langs/yaml.mjs"),
};

let highlighterPromise: Promise<HighlighterCore> | null = null;
const languagePromises = new Map<string, Promise<void>>();

async function getCodeHighlighter() {
  if (!highlighterPromise) {
    highlighterPromise = Promise.all([
      import("shiki/core"),
      import("shiki/engine/javascript"),
    ]).then(([core, engine]) =>
      core.createHighlighterCore({
        themes: [
          core.createCssVariablesTheme({
            name: "yurou-code",
            variablePrefix: "--code-",
          }),
        ],
        langs: [],
        engine: engine.createJavaScriptRegexEngine(),
      }),
    );
  }

  return highlighterPromise;
}

async function loadCodeLanguage(highlighter: HighlighterCore, language: string) {
  if (highlighter.getLoadedLanguages().includes(language)) return;

  let promise = languagePromises.get(language);
  if (!promise) {
    const loader = languageLoaders[language];
    if (!loader) return;
    promise = loader().then((module) => highlighter.loadLanguage(module.default));
    languagePromises.set(language, promise);
  }

  await promise;
}

export async function highlightCodeToHtml(code: string, language: string) {
  const highlighter = await getCodeHighlighter();
  await loadCodeLanguage(highlighter, language);

  if (!highlighter.getLoadedLanguages().includes(language)) return null;

  return highlighter.codeToHtml(code, {
    lang: language,
    theme: "yurou-code",
  });
}

export async function highlightCodeToTokens(
  code: string,
  language: string,
): Promise<ThemedToken[][] | null> {
  const highlighter = await getCodeHighlighter();
  await loadCodeLanguage(highlighter, language);

  if (!highlighter.getLoadedLanguages().includes(language)) return null;

  return highlighter.codeToTokens(code, {
    lang: language,
    theme: "yurou-code",
  }).tokens;
}
