import MarkdownIt from 'markdown-it'
import mdKatex from '@traptitech/markdown-it-katex'
import {
  createHighlighter,
  type BundledLanguage,
  type Highlighter,
} from 'shiki'

let highlighter: Highlighter | null = null

const allowed_langs: BundledLanguage[] = [
  'python',
  'javascript',
  'typescript',
  'typst',
  'markdown',
  'json',
  'toml',
  'yaml',
  'bash',
  'c',
  'c++',
  'rust',
  'go',
  'zig',
  'makefile',
  'make',
  'nim',
  'nix',
  'kdl',
  'md',
  'sh',
  'mermaid',
  'cpp',
  'fish',
  'zsh',
  'docker',
  'dockerfile',
  'yml',
  'system-verilog',
  'sql',
  'glsl',
  'mediawiki',
  'systemd',
  'ini',
]

export async function ensureShikiEngine() {
  if (highlighter === null) {
    highlighter = await createHighlighter({
      themes: ['one-light', 'one-dark-pro'],
      /* 额...要我自己定义需要的所有语言吗 */
      langs: allowed_langs,
    })
  }
  return highlighter
}

export function shikiRender(code: string, lang: string) {
  return highlighter!.codeToHtml(code, {
    lang: (allowed_langs as string[]).indexOf(lang) != -1 ? lang : 'text',
    themes: {
      light: 'one-light',
      dark: 'one-dark-pro',
    },
    defaultColor: false,
  })
}

/**
 * KaTeX 渲染结果修改。
 *
 * 让 KaTeX 的公式可复制。
 */

function setupKatexCopy(md: MarkdownIt) {
  const defaultInline = md.renderer.rules.math_inline!
  const defaultBlock = md.renderer.rules.math_block!

  md.renderer.rules.math_inline = (tokens, idx, options, env, self) => {
    const content = tokens[idx].content
    const rendered = defaultInline(tokens, idx, options, env, self)

    return `<span class="katex-wrapper inline" data-code="${encodeURIComponent(content)}">${rendered}</span>`
  }

  md.renderer.rules.math_block = (tokens, idx, options, env, self) => {
    const content = tokens[idx].content
    const rendered = defaultBlock(tokens, idx, options, env, self)
    return `
      <div class="katex-wrapper block" data-code="${encodeURIComponent(content)}">${rendered}
        <button class="copy-katex" title="Copy LaTeX">
复制
        </button>
      </div>`
  }
}

/**
 * Markdown 渲染函数。
 *
 * 为了让客户端和服务端表现相同，所以单独拿出来了一个 Markdown 渲染器。
 */
export async function renderMarkdown(content: string): Promise<string> {
  await ensureShikiEngine()

  const md = new MarkdownIt({
    html: true,
    linkify: true,
    highlight: shikiRender,
  })

  md.use(mdKatex)
  setupKatexCopy(md)

  return md.render(content)
}
