import katex from 'katex'
import type { MarkedExtension, TokenizerAndRendererExtension, Tokens } from 'marked'

/**
 * marked 拡張として数式 (KaTeX) を描画する。
 *
 * 対応する区切り記号:
 *   - ブロック: `$$ ... $$` / `\[ ... \]`（単独行）
 *   - インライン: `$ ... $` / `\( ... \)` / 行中の `$$ ... $$`
 *
 * `$` はドル記号としても書かれるため、インラインの `$ ... $` だけは
 * 誤爆を避けるための条件を付けている（下記 INLINE_DOLLAR 参照）。
 * 明示的にドル記号を書きたい場合は Markdown のエスケープ `\$` が使える。
 */

// ブロック: 行頭（インデント 3 まで）から始まる $$ ... $$ / \[ ... \]
const BLOCK_DOLLAR = /^ {0,3}\$\$([\s\S]+?)\$\$[ \t]*(?:\n+|$)/
const BLOCK_BRACKET = /^ {0,3}\\\[([\s\S]+?)\\\][ \t]*(?:\n+|$)/

// インライン: \( ... \) と、段落の途中に置かれた $$ ... $$
const INLINE_PAREN = /^\\\(((?:\\[^]|[^\\])+?)\\\)/
const INLINE_DOUBLE = /^\$\$((?:\\[^]|[^\\])+?)\$\$/
// インライン: $ ... $
//   - 開き `$` の直後に空白・`$` が来るものは対象外
//   - 中身に改行と裸の `$` を含まない
//   - 閉じ `$` の直後が数字のものは対象外（"$100から$200" のような金額表記の除外）
//   - 中身が空白で終わるものは対象外（tokenizer 側で判定）
const INLINE_DOLLAR = /^\$(?![\s$])((?:\\.|[^\\$\n])+?)\$(?!\d)/

function escapeHtml(src: string): string {
  return src
    .replace(/&/g, '&')
    .replace(/</g, '&lt;')
    .replace(/>/g, '&gt;')
    .replace(/"/g, '&quot;')
}

function renderMath(tex: string, displayMode: boolean, raw: string): string {
  try {
    return katex.renderToString(tex.trim(), {
      displayMode,
      throwOnError: false, // 構文エラーは例外にせず赤字で表示する
      errorColor: '#f87171', // ダーク背景で読める色にする
      strict: false, // 日本語混じりなどの警告を出さない
      output: 'htmlAndMathml',
    })
  } catch {
    // KaTeX が想定外の失敗をした場合は元の文字列をそのまま見せる
    return escapeHtml(raw)
  }
}

const blockMath: TokenizerAndRendererExtension = {
  name: 'blockMath',
  level: 'block',
  start(src: string) {
    return src.match(/(?:^|\n) {0,3}(?:\$\$|\\\[)/)?.index
  },
  tokenizer(src: string) {
    const match = BLOCK_DOLLAR.exec(src) ?? BLOCK_BRACKET.exec(src)
    if (!match) return undefined
    return { type: 'blockMath', raw: match[0], text: match[1] }
  },
  renderer(token: Tokens.Generic) {
    return renderMath(token.text as string, true, token.raw as string)
  },
}

const inlineMath: TokenizerAndRendererExtension = {
  name: 'inlineMath',
  level: 'inline',
  start(src: string) {
    // エスケープされた \$ は数式の開始とみなさない
    const match = /(^|[^\\])(\$|\\\()/.exec(src)
    return match ? match.index + match[1].length : undefined
  },
  tokenizer(src: string) {
    const paren = INLINE_PAREN.exec(src)
    if (paren) return { type: 'inlineMath', raw: paren[0], text: paren[1], display: false }

    const double = INLINE_DOUBLE.exec(src)
    if (double) return { type: 'inlineMath', raw: double[0], text: double[1], display: true }

    const single = INLINE_DOLLAR.exec(src)
    if (single && !/[\s\\]$/.test(single[1])) {
      return { type: 'inlineMath', raw: single[0], text: single[1], display: false }
    }
    return undefined
  },
  renderer(token: Tokens.Generic) {
    return renderMath(token.text as string, token.display === true, token.raw as string)
  },
}

export const mathExtension: MarkedExtension = {
  extensions: [blockMath, inlineMath],
}

/** KaTeX が出力する MathML 用に DOMPurify へ追加で許可するタグ・属性 */
export const mathSanitizeConfig = {
  ADD_TAGS: ['semantics', 'annotation'],
  ADD_ATTR: ['encoding'],
}