import MermaidBlock from './MermaidBlock'
import React from 'react'
import { codeToHtml } from 'shiki'

function extractText(node: React.ReactNode): string {
  if (typeof node === 'string') return node
  if (typeof node === 'number') return String(node)
  if (Array.isArray(node)) return node.map(extractText).join('')
  if (React.isValidElement(node)) {
    const props = node.props as { children?: React.ReactNode }
    return extractText(props.children)
  }
  return ''
}

function extractLang(className?: string): string {
  const match = className?.match(/language-(\w+)/)
  return match ? match[1] : 'text'
}

async function HighlightedCode({ code, lang }: { code: string; lang: string }) {
  try {
    const html = await codeToHtml(code, {
      lang,
      theme: 'github-dark-default',
    })
    return (
      <div
        className="not-prose my-6 overflow-x-auto rounded-xl text-sm shadow-sm [&>pre]:p-5 [&>pre]:leading-relaxed"
        dangerouslySetInnerHTML={{ __html: html }}
      />
    )
  } catch {
    return <pre className="not-prose my-6 overflow-x-auto rounded-xl bg-gray-950 text-gray-100 p-5 text-sm leading-relaxed shadow-sm">{code}</pre>
  }
}

export default function PreBlock({
  children,
  ...rest
}: React.HTMLAttributes<HTMLPreElement> & { children?: React.ReactNode }) {
  if (React.isValidElement(children)) {
    const childProps = children.props as { className?: string; children?: React.ReactNode }
    const code = extractText(childProps.children)
    const lang = extractLang(childProps.className)

    if (lang === 'mermaid') {
      return <MermaidBlock code={code} />
    }

    return <HighlightedCode code={code.trimEnd()} lang={lang} />
  }
  return <pre {...rest}>{children}</pre>
}
