jiang13-bbs/frontend/components/markdown/shikiHighlighter.ts

import { createHighlighterCore, type HighlighterCore } from "shiki/core";
import { createOnigurumaEngine } from "shiki/engine/oniguruma";
import type { LanguageRegistration } from "shiki";
import { aardioGrammar } from "./aardioGrammar";

/** 站点 Markdown 代码高亮：显式白名单（非全量） */
export const SHIKI_LANG_IDS = [
  "go",
  "python",
  "r",
  "markdown",
  "javascript",
  "typescript",
  "tsx",
  "jsx",
  "json",
  "jsonc",
  "yaml",
  "toml",
  "ini",
  "bash",
  "sql",
  "html",
  "css",
  "scss",
  "diff",
  "rust",
  "ruby",
  "powershell",
  "nginx",
  "dockerfile",
  "kotlin",
  "java",
  "c",
  "cpp",
  "csharp",
  "php",
  "swift",
  "dart",
  "lua",
  "vue",
  "graphql",
  "xml",
  "regex",
  "makefile",
  "prisma",
  "plaintext",
] as const;

export type SiteShikiLang = (typeof SHIKI_LANG_IDS)[number] | "aardio";

export const SHIKI_LANG_ALIAS: Record<string, string> = {
  py: "python",
  md: "markdown",
  js: "javascript",
  ts: "typescript",
  sh: "bash",
  shell: "bash",
  zsh: "bash",
  yml: "yaml",
  text: "plaintext",
  txt: "plaintext",
  golang: "go",
  rb: "ruby",
  ps1: "powershell",
  pwsh: "powershell",
  "c++": "cpp",
  cs: "csharp",
  "c#": "csharp",
  regexp: "regex",
  docker: "dockerfile",
};

/**
 * 白名单 id → @shikijs/langs grammar 内部注册名。
 * 部分语言的 dist 文件只是别名重导出（如 bash.mjs = shellscript.mjs），
 * HighlighterCore 按内部名注册，codeToHtml(lang:"bash") 会静默降级 plaintext。
 */
const INTERNAL_LANG_ALIAS: Record<string, string> = {
  bash: "shellscript",
  sh: "shellscript",
  shell: "shellscript",
  zsh: "shellscript",
  dockerfile: "docker",
  regex: "regexp",
  makefile: "make",
};

/**
 * createHighlighterCore 的 langAlias：输入别名 + 白名单 id 都要能落到 grammar 内部名。
 * 注意 core 会校验别名链无环，且 docker/regexp 本身就是内部注册名，
 * 不能再作为 alias 键（会与 dockerfile→docker、regex→regexp 互指成环）。
 */
function buildCoreLangAlias(): Record<string, string> {
  const merged = { ...SHIKI_LANG_ALIAS, ...INTERNAL_LANG_ALIAS };
  delete merged["docker"];
  delete merged["regexp"];
  return merged;
}

export const SHIKI_THEMES = {
  light: "github-light",
  dark: "github-dark-dimmed",
} as const;

const ALLOWED = new Set<string>([...SHIKI_LANG_IDS, "aardio"]);

/** 未知语言降级为 plaintext；别名归一到白名单 id */
export function resolveShikiLang(raw?: string | null): SiteShikiLang {
  const id = (raw || "").trim().toLowerCase();
  if (!id) return "plaintext";
  const mapped = SHIKI_LANG_ALIAS[id] || id;
  if (ALLOWED.has(mapped)) return mapped as SiteShikiLang;
  return "plaintext";
}

let highlighterPromise: Promise<HighlighterCore> | null = null;

/** 站点单例 highlighter：仅白名单语言 + aardio */
export function getSiteHighlighter(): Promise<HighlighterCore> {
  if (!highlighterPromise) {
    highlighterPromise = createHighlighterCore({
      themes: [
        import("@shikijs/themes/github-light"),
        import("@shikijs/themes/github-dark-dimmed"),
      ],
      langs: [
        import("@shikijs/langs/go"),
        import("@shikijs/langs/python"),
        import("@shikijs/langs/r"),
        import("@shikijs/langs/markdown"),
        import("@shikijs/langs/javascript"),
        import("@shikijs/langs/typescript"),
        import("@shikijs/langs/tsx"),
        import("@shikijs/langs/jsx"),
        import("@shikijs/langs/json"),
        import("@shikijs/langs/yaml"),
        import("@shikijs/langs/toml"),
        import("@shikijs/langs/bash"),
        import("@shikijs/langs/sql"),
        import("@shikijs/langs/html"),
        import("@shikijs/langs/css"),
        import("@shikijs/langs/scss"),
        import("@shikijs/langs/diff"),
        import("@shikijs/langs/rust"),
        import("@shikijs/langs/ruby"),
        import("@shikijs/langs/powershell"),
        import("@shikijs/langs/nginx"),
        import("@shikijs/langs/dockerfile"),
        import("@shikijs/langs/kotlin"),
        import("@shikijs/langs/java"),
        import("@shikijs/langs/c"),
        import("@shikijs/langs/cpp"),
        import("@shikijs/langs/csharp"),
        import("@shikijs/langs/php"),
        import("@shikijs/langs/swift"),
        import("@shikijs/langs/dart"),
        import("@shikijs/langs/lua"),
        import("@shikijs/langs/vue"),
        import("@shikijs/langs/graphql"),
        import("@shikijs/langs/xml"),
        import("@shikijs/langs/regex"),
        import("@shikijs/langs/makefile"),
        import("@shikijs/langs/prisma"),
        // plaintext / text / txt 为 Shiki 特殊语言，无需 grammar
        aardioGrammar,
      ],
      langAlias: buildCoreLangAlias(),
      engine: createOnigurumaEngine(import("shiki/wasm")),
    });
  }
  return highlighterPromise;
}

/** 客户端/工具：按白名单高亮为双主题 HTML */
export async function siteCodeToHtml(code: string, langRaw?: string | null): Promise<string> {
  const highlighter = await getSiteHighlighter();
  const lang = resolveShikiLang(langRaw);
  return highlighter.codeToHtml(code, {
    lang,
    themes: {
      light: SHIKI_THEMES.light,
      dark: SHIKI_THEMES.dark,
    },
    defaultColor: false,
  });
}

/**
 * rehype-pretty-code 会对每个 fence 调用 loadLanguage(langId: string)。
 * HighlighterCore 不能把裸 string 当 grammar（会 TypeError 并可能搞坏注册表），
 * 这里包装：已装载的语言直接跳过。
 */
export async function getPrettyCodeHighlighter(): Promise<HighlighterCore> {
  const core = await getSiteHighlighter();

  const loadLanguage = async (
    ...langs: (string | LanguageRegistration | Promise<LanguageRegistration>)[]
  ) => {
    const pending: (LanguageRegistration | Promise<LanguageRegistration>)[] = [];
    const loaded = new Set(core.getLoadedLanguages());
    for (const lang of langs.flat()) {
      if (typeof lang === "string") {
        const id = resolveShikiLang(lang);
        if (loaded.has(lang) || loaded.has(id) || id === "plaintext") continue;
        continue;
      }
      pending.push(lang);
    }
    if (pending.length > 0) {
      await core.loadLanguage(...pending);
    }
  };

  return new Proxy(core, {
    get(target, prop, receiver) {
      if (prop === "loadLanguage") return loadLanguage;
      const value = Reflect.get(target, prop, receiver);
      return typeof value === "function" ? value.bind(target) : value;
    },
  }) as HighlighterCore;
}