import { marked } from 'marked';
import Prism from 'prismjs';
import renderMathInElement from 'katex/contrib/auto-render';

let latexRenderTimer: number | null = null;
const markdownCache = new Map<string, string>();

function wrapCodeBlocks(html: string, isStreaming = false) {
  if (isStreaming) {
    return html;
  }

  let counter = 0;
  return html.replace(/<pre><code([^>]*)>([\s\S]*?)<\/code><\/pre>/g, (match, attributes, content) => {
    const langMatch = attributes.match(/class="[^"]*language-(\w+)/);
    const language = langMatch ? langMatch[1] : 'text';
    const blockId = `code-${Date.now()}-${counter++}`;
    const escapedContent = content
      .replace(/&/g, '&amp;')
      .replace(/</g, '&lt;')
      .replace(/>/g, '&gt;')
      .replace(/"/g, '&quot;');

    return `
            <div class="code-block-wrapper">
                <div class="code-block-header">
                    <span class="code-language">${language}</span>
                    <button class="copy-code-btn" data-code="${blockId}" title="复制代码" aria-label="复制代码"></button>
                </div>
                <pre><code${attributes} data-code-id="${blockId}" data-original-code="${escapedContent}">${content}</code></pre>
            </div>`;
  });
}

export function renderMarkdown(text: string, isStreaming = false) {
  if (!text) {
    return '';
  }

  if (typeof marked === 'undefined') {
    return text;
  }

  marked.setOptions({
    breaks: true,
    gfm: true
  });

  if (!isStreaming) {
    const cacheKey = `${text.length}_${text.substring(0, 100)}`;
    if (markdownCache.has(cacheKey)) {
      return markdownCache.get(cacheKey) as string;
    }
  }

  const parsed = marked.parse(text) as string;
  let html = parsed;
  html = wrapCodeBlocks(html, isStreaming);

  if (!isStreaming && text.length < 10000) {
    const cacheKey = `${text.length}_${text.substring(0, 100)}`;
    markdownCache.set(cacheKey, html);
    if (markdownCache.size > 20) {
      const firstKey = markdownCache.keys().next().value as string | undefined;
      if (firstKey) {
        markdownCache.delete(firstKey);
      }
    }
  }

  if (!isStreaming) {
    setTimeout(() => {
      if (typeof Prism !== 'undefined') {
        const codeBlocks = document.querySelectorAll('.code-block-wrapper pre code:not([data-highlighted])');
        codeBlocks.forEach(block => {
          try {
            Prism.highlightElement(block as HTMLElement);
            block.setAttribute('data-highlighted', 'true');
          } catch (error) {
            console.warn('代码高亮失败:', error);
          }
        });
      }

      if (typeof renderMathInElement !== 'undefined') {
        const elements = document.querySelectorAll('.text-output .text-content:not(.streaming-text)');
        elements.forEach(element => {
          if (element.hasAttribute('data-math-rendered')) return;
          try {
            renderMathInElement(element as HTMLElement, {
              delimiters: [
                { left: '$$', right: '$$', display: true },
                { left: '$', right: '$', display: false },
                { left: '\\[', right: '\\]', display: true },
                { left: '\\(', right: '\\)', display: false }
              ],
              throwOnError: false,
              trust: true
            });
            element.setAttribute('data-math-rendered', 'true');
          } catch (error) {
            console.warn('LaTeX渲染失败:', error);
          }
        });
      }
    }, 100);
  }

  return html;
}

export function renderLatexInRealtime() {
  if (typeof renderMathInElement === 'undefined') {
    return;
  }

  if (latexRenderTimer) {
    cancelAnimationFrame(latexRenderTimer);
  }

  latexRenderTimer = requestAnimationFrame(() => {
    const elements = document.querySelectorAll('.text-output .streaming-text');
    elements.forEach(element => {
      try {
        renderMathInElement(element as HTMLElement, {
          delimiters: [
            { left: '$$', right: '$$', display: true },
            { left: '$', right: '$', display: false },
            { left: '\\[', right: '\\]', display: true },
            { left: '\\(', right: '\\)', display: false }
          ],
          throwOnError: false,
          trust: true
        });
      } catch {
        // ignore
      }
    });
  });
}
