"use client";

import React, { useState, useMemo } from 'react';
import ReactMarkdown from 'react-markdown';
import remarkGfm from 'remark-gfm';
import remarkMath from 'remark-math';
import rehypeKatex from 'rehype-katex';
import rehypeRaw from 'rehype-raw';
import rehypeSlug from 'rehype-slug';
import rehypeAutolinkHeadings from 'rehype-autolink-headings';
import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
import { vscDarkPlus, coy } from 'react-syntax-highlighter/dist/cjs/styles/prism';
import 'katex/dist/katex.min.css';
import { useTheme } from 'next-themes';

// --- TYPE DEFINITIONS ---
interface MarkdownRendererProps {
  content: string;
  className?: string;
}

type MarkdownComponentProps = {
  node?: any;
  children?: React.ReactNode;
  className?: string;
  [key: string]: any;
};

// --- PRE-PROCESSOR ---
function preprocessMarkdown(content: string): string {
  if (!content) {
    return "";
  }
  let processed = content.replace(/\[cite:\s*(\d+(?:,\s*\d+)*)\]/g, '<span class="citation" title="Citation reference">[$1]</span>');
  return processed;
}

// --- HELPER: Copy to Clipboard ---
const copyToClipboard = (text: string, onCopy: () => void) => {
  if (navigator.clipboard) {
    navigator.clipboard.writeText(text).then(onCopy).catch(err => {
      console.error('Async clipboard copy failed:', err);
    });
  } else {
    const textArea = document.createElement('textarea');
    textArea.value = text;
    textArea.style.position = 'fixed';
    textArea.style.opacity = '0';
    document.body.appendChild(textArea);
    textArea.focus();
    textArea.select();
    try {
      document.execCommand('copy');
      onCopy();
    } catch (err) {
      console.error('Fallback execCommand copy failed:', err);
    }
    document.body.removeChild(textArea);
  }
};


// --- CUSTOM RENDERER COMPONENTS ---

const CodeBlock: React.FC<MarkdownComponentProps> = React.memo(({ node, inline, className, children }) => {
  const { theme } = useTheme();
  const [isCopied, setIsCopied] = useState(false);
  const match = /language-(\w+)/.exec(className || '');
  const lang = match ? match[1] : 'text';
  const codeString = String(children).replace(/\n$/, '');

  const handleCopy = () => {
    copyToClipboard(codeString, () => {
      setIsCopied(true);
      setTimeout(() => setIsCopied(false), 2000);
    });
  };

  if (inline) {
    return (
      <code className="py-0.5 px-1.5 rounded text-sm bg-gray-100 dark:bg-gray-800 text-indigo-600 dark:text-indigo-400 font-mono">
        {children}
      </code>
    );
  }

  return (
    <div className="code-block-container relative my-6 rounded-lg shadow-lg bg-gray-900 dark:bg-black/50 border border-gray-200 dark:border-gray-700">
      <div className="code-header flex items-center justify-between px-4 py-2 bg-gray-50 dark:bg-gray-800/50 border-b border-gray-200 dark:border-gray-700">
        <span className="text-xs font-medium text-gray-500 dark:text-gray-400 uppercase tracking-wider">
          {lang}
        </span>
        <button
          onClick={handleCopy}
          className={`copy-button flex items-center gap-1.5 text-xs py-1 px-2 rounded-md transition-all duration-200 ${
            isCopied
              ? 'bg-green-100 dark:bg-green-900/30 text-green-600 dark:text-green-400'
              : 'bg-gray-100 dark:bg-gray-700 text-gray-600 dark:text-gray-300 hover:bg-indigo-100 dark:hover:bg-indigo-900/50 hover:text-indigo-500 dark:hover:text-indigo-400'
          }`}
          title="Copy to clipboard"
        >
          {isCopied ? (
            <>
              <svg xmlns="http://www.w3.org/2000/svg" className="h-3.5 w-3.5" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round"><path d="M20 6L9 17l-5-5" /></svg>
              <span>Copied!</span>
            </>
          ) : (
            <>
              <svg xmlns="http://www.w3.org/2000/svg" className="h-3.5 w-3.5" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round"><rect x="9" y="9" width="13" height="13" rx="2" ry="2" /><path d="M5 15V5a2 2 0 012-2h8" /></svg>
              <span>Copy</span>
            </>
          )}
        </button>
      </div>
      <SyntaxHighlighter
        language={lang}
        style={theme === 'dark' ? vscDarkPlus : coy}
        showLineNumbers={codeString.split('\n').length > 1}
        wrapLines={true}
        customStyle={{
          margin: 0,
          padding: '1rem',
          background: 'transparent',
          fontSize: '0.9rem',
        }}
        codeTagProps={{
          style: {
            fontFamily: '"JetBrains Mono", "Fira Code", monospace',
            lineHeight: 1.6,
          },
        }}
        lineNumberStyle={{
          minWidth: '2.5em',
          paddingRight: '1em',
          textAlign: 'right',
          userSelect: 'none',
          color: theme === 'dark' ? '#6b7280' : '#9ca3af',
        }}
      >
        {codeString}
      </SyntaxHighlighter>
    </div>
  );
});
CodeBlock.displayName = 'CodeBlock';

const CustomImage: React.FC<React.ImgHTMLAttributes<HTMLImageElement>> = React.memo(({ src, alt, title, ...props }) => {
    const [isZoomed, setIsZoomed] = useState(false);
    return (
        <div className="my-6 text-center">
            <img 
                src={src}
                alt={alt || 'Markdown Image'}
                title={title}
                className="max-w-full h-auto rounded-lg shadow-md hover:shadow-xl transition-all duration-300 cursor-zoom-in mx-auto block"
                onClick={() => setIsZoomed(true)}
                loading="lazy"
                {...props}
            />
            {title && <p className="text-sm text-gray-500 dark:text-gray-400 mt-2 italic">{title}</p>}
            {isZoomed && (
                <div 
                    className="fixed inset-0 z-50 bg-black bg-opacity-80 flex items-center justify-center p-4 cursor-zoom-out animate-fade-in"
                    onClick={() => setIsZoomed(false)}
                >
                    <img src={src} alt={alt || ''} className="max-w-[95vw] max-h-[95vh] object-contain rounded-lg" />
                </div>
            )}
        </div>
    );
});
CustomImage.displayName = 'CustomImage';


// --- MAIN RENDERER COMPONENT ---
const MarkdownRenderer: React.FC<MarkdownRendererProps> = ({ content, className = '' }) => {
  const processedContent = useMemo(() => preprocessMarkdown(content), [content]);

  const components = useMemo(() => ({
    pre: ({ node, ...props }: MarkdownComponentProps) => {
        const childrenArray = React.Children.toArray(props.children);
        const child = childrenArray[0];
        if (React.isValidElement(child) && (child.props as any)?.className?.includes('language-')) {
            return <CodeBlock node={node} {...(child.props as any)} />;
        }
        return <pre className="p-4 my-4 bg-gray-100 dark:bg-gray-800 rounded-lg overflow-x-auto" {...props} />;
    },
    code: ({ node, inline, className, children, ...props }: MarkdownComponentProps) => {
        return <CodeBlock node={node} inline={inline} className={className} {...props}>{children}</CodeBlock>;
    },
    div: ({ className, children, ...props }: MarkdownComponentProps) => {
        if (className?.includes('math-display')) {
            return <div className="my-6 overflow-x-auto text-center">{children}</div>;
        }
        return <div className={className} {...props}>{children}</div>;
    },
    img: CustomImage,
    h1: ({node, ...props}: MarkdownComponentProps) => <h1 className="text-4xl font-extrabold mt-10 mb-6 pb-3 border-b-2 border-gray-200 dark:border-gray-700" {...props} />,
    h2: ({node, ...props}: MarkdownComponentProps) => <h2 className="text-3xl font-bold mt-8 mb-5 pb-2 border-b border-gray-200 dark:border-gray-700" {...props} />,
    h3: ({node, ...props}: MarkdownComponentProps) => <h3 className="text-2xl font-semibold mt-7 mb-4" {...props} />,
    h4: ({node, ...props}: MarkdownComponentProps) => <h4 className="text-xl font-semibold mt-6 mb-3" {...props} />,
    p: ({node, ...props}: MarkdownComponentProps) => <p className="leading-relaxed my-4" {...props} />,
    a: ({node, ...props}: MarkdownComponentProps) => <a className="text-indigo-600 dark:text-indigo-400 hover:underline decoration-2 underline-offset-2" target="_blank" rel="noopener noreferrer" {...props} />,
    blockquote: ({node, ...props}: MarkdownComponentProps) => <blockquote className="pl-5 py-2 my-5 border-l-4 border-indigo-500 bg-indigo-50 dark:bg-indigo-900/20 text-gray-700 dark:text-gray-300" {...props} />,
    ul: ({node, ...props}: MarkdownComponentProps) => <ul className="list-disc list-outside pl-6 my-4 space-y-1.5 academic-list" {...props} />,
    ol: ({node, ...props}: MarkdownComponentProps) => <ol className="list-decimal list-outside pl-6 my-4 space-y-1.5 academic-list" {...props} />,
    li: ({node, ...props}: MarkdownComponentProps) => <li className="pl-2 leading-relaxed academic-list-item" {...props} />,
    table: ({node, ...props}: MarkdownComponentProps) => (
      <div className="my-6 overflow-x-auto rounded-lg border border-gray-200 dark:border-gray-700 shadow-sm">
        <table className="min-w-full divide-y divide-gray-200 dark:divide-gray-700 academic-table" {...props} />
      </div>
    ),
    thead: ({node, ...props}: MarkdownComponentProps) => <thead className="bg-gradient-to-r from-blue-50 to-indigo-50 dark:from-gray-800 dark:to-gray-700" {...props} />,
    th: ({node, ...props}: MarkdownComponentProps) => (
      <th className="px-4 py-3 text-left text-xs font-semibold text-gray-700 dark:text-gray-300 uppercase tracking-wider border-b-2 border-gray-200 dark:border-gray-600" {...props} />
    ),
    td: ({node, ...props}: MarkdownComponentProps) => (
      <td className="px-4 py-3 text-sm text-gray-700 dark:text-gray-300 border-b border-gray-100 dark:border-gray-700 academic-table-cell" {...props} />
    ),
    hr: ({node, ...props}: MarkdownComponentProps) => <hr className="my-8 border-gray-200 dark:border-gray-700" {...props} />,
  }), []);

  const remarkPlugins = [remarkGfm, remarkMath];
  const rehypePlugins: any[] = [
    rehypeRaw,
    rehypeSlug,
    [rehypeAutolinkHeadings, { behavior: 'wrap' }],
    [rehypeKatex, { 
      throwOnError: false, 
      errorColor: '#EF4444',
      strict: false,
      output: 'html',
      macros: {
        "\\vec": "\\overrightarrow{#1}",
        "\\hat": "\\hat{#1}",
        "\\perp": "\\perp",
        "\\ihat": "\\hat{\\imath}",
        "\\jhat": "\\hat{\\jmath}",
        "\\khat": "\\hat{k}",
        "\\R": "\\mathbb{R}",
        "\\N": "\\mathbb{N}",
        "\\Z": "\\mathbb{Z}",
        "\\Q": "\\mathbb{Q}",
        "\\C": "\\mathbb{C}"
      }
    }],
  ];

  return (
    <>
      <style jsx global>{`
        /* --- KaTeX Math Styling --- */
        .katex-display {
          margin: 1em 0 !important;
          text-align: center !important;
          overflow-x: auto;
          overflow-y: hidden;
          padding: 0.5em 0;
        }
        .katex {
          font-size: 1.1em !important;
          line-height: 1.2;
        }
        .katex-display > .katex {
          font-size: 1.25em !important;
        }
        .katex .text {
          font-family: "KaTeX_Main", "Times New Roman", serif;
        }
        .katex-error {
          color: #dc2626;
          background-color: #fef2f2;
          border: 1px solid #fecaca;
          padding: 2px 6px;
          border-radius: 4px;
        }
        .dark .katex-error {
          color: #f87171;
          background-color: #2d1b1b;
          border-color: #7f1d1d;
        }
        .katex .vlist-t {
          vertical-align: baseline;
        }
        .katex-display .katex > .katex-html {
          max-width: 100%;
          overflow-x: auto;
        }

        /* --- Custom Citation Styling --- */
        .prose .citation {
          font-size: 0.8em;
          color: #4f46e5;
          background-color: #eef2ff;
          padding: 2px 6px;
          border-radius: 4px;
          margin: 0 2px;
          font-weight: 500;
          border: 1px solid #c7d2fe;
          display: inline-block;
          text-decoration: none;
          transition: all 0.2s ease;
        }
        .prose .citation:hover {
          background-color: #e0e7ff;
          transform: translateY(-1px);
        }
        .dark .prose .citation {
          color: #a5b4fc;
          background-color: #312e81;
          border-color: #4f46e5;
        }
        .dark .prose .citation:hover {
          background-color: #3730a3;
        }

        /* --- General & Animation --- */
        @keyframes fade-in {
          from { opacity: 0; }
          to { opacity: 1; }
        }
        .animate-fade-in {
          animation: fade-in 0.2s ease-out;
        }

        /* --- Academic Content Styling --- */
        .academic-table {
          font-family: system-ui, -apple-system, sans-serif;
          border-collapse: separate;
          border-spacing: 0;
        }
        .academic-table-cell {
          vertical-align: top;
          line-height: 1.6;
        }
        .academic-table tbody tr:hover {
          background-color: #f8fafc;
          transition: background-color 0.2s ease;
        }
        .dark .academic-table tbody tr:hover {
          background-color: #374151;
        }
        .academic-list-item {
          margin: 0.5rem 0;
          padding-left: 0.5rem;
        }
        .prose ul li::marker, .prose ol li::marker {
          color: #6366f1;
          font-weight: bold;
        }
        .dark .prose ul li::marker, .dark .prose ol li::marker {
          color: #818cf8;
        }

        /* --- Heading Enhancements --- */
        .prose h2 {
          color: #1e40af;
          position: relative;
        }
        .prose h2::after {
          content: '';
          position: absolute;
          bottom: -4px;
          left: 0;
          width: 60px;
          height: 3px;
          background: linear-gradient(90deg, #3b82f6, #8b5cf6);
          border-radius: 2px;
        }
        .dark .prose h2 {
          color: #60a5fa;
        }
        .prose h3 {
          color: #1d4ed8;
          border-left: 4px solid #3b82f6;
          padding-left: 12px;
          margin-left: -16px;
        }
        .dark .prose h3 {
          color: #93c5fd;
          border-left-color: #60a5fa;
        }
        
        /* --- Responsive & Print --- */
        @media (max-width: 768px) {
          .academic-table { font-size: 0.875rem; }
          .prose h1 { font-size: 2rem; }
          .prose h2 { font-size: 1.5rem; }
          .prose h3 { font-size: 1.25rem; }
        }
        @media print {
          .academic-table { border-collapse: collapse; border: 1px solid #000; }
          .academic-table th, .academic-table td { border: 1px solid #000; padding: 4px 8px; }
        }
      `}</style>
      
      <div className={`markdown-renderer prose dark:prose-invert prose-indigo max-w-none ${className}`}>
        <ReactMarkdown
          remarkPlugins={remarkPlugins}
          rehypePlugins={rehypePlugins}
          components={components}
        >
          {processedContent}
        </ReactMarkdown>
      </div>
    </>
  );
};

export default React.memo(MarkdownRenderer);