import type { MDXComponents } from 'mdx/types'
import { isValidElement } from 'react'
import Link from 'next/link'
import { ClickableHeading } from '@/app/components/docs/ClickableHeading'
import { DocsPageMenu } from '@/app/components/docs/DocsPageMenu'
import { Updated } from '@/app/components/docs/Updated'

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

export function useMDXComponents(components: MDXComponents): MDXComponents {
    return {
        h1: ({ children }) => (
            <div className="flex items-start justify-between gap-2">
                <ClickableHeading level={1} text={extractHeadingText(children)}>{children}</ClickableHeading>
                <DocsPageMenu />
            </div>
        ),
        h2: ({ children }) => (
            <ClickableHeading level={2} text={extractHeadingText(children)}>{children}</ClickableHeading>
        ),
        h3: ({ children }) => (
            <ClickableHeading level={3} text={extractHeadingText(children)}>{children}</ClickableHeading>
        ),
        h4: ({ children }) => (
            <ClickableHeading level={4} text={extractHeadingText(children)}>{children}</ClickableHeading>
        ),
        p: ({ children }) => (
            <p className="text-base sm:text-lg leading-relaxed tracking-tight text-muted mb-4">
                {children}
            </p>
        ),
        a: ({ href, children }) => {
            const isInternal = href && (href.startsWith('/') || href.startsWith('#'))
            const className = "text-base-content underline underline-offset-2 hover:text-muted transition-colors"

            if (isInternal) {
                return (
                    <Link href={href} className={className}>
                        {children}
                    </Link>
                )
            }

            return (
                <a href={href} className={className} target="_blank" rel="noopener noreferrer">
                    {children}
                </a>
            )
        },
        strong: ({ children }) => (
            <strong className="font-medium text-base-content">{children}</strong>
        ),
        code: ({ children }) => (
            <code className="font-mono text-sm bg-base-200 px-1.5 py-0.5 rounded text-base-content">
                {children}
            </code>
        ),
        pre: ({ children }) => (
            <pre className="font-mono text-sm bg-base-200 border border-border rounded-xl p-4 overflow-x-auto mb-6">
                {children}
            </pre>
        ),
        ul: ({ children }) => (
            <ul className="list-disc list-outside ml-5 text-muted mb-4 space-y-2">
                {children}
            </ul>
        ),
        ol: ({ children }) => (
            <ol className="list-decimal list-outside ml-5 text-muted mb-4 space-y-2">
                {children}
            </ol>
        ),
        li: ({ children }) => (
            <li className="text-base sm:text-lg leading-relaxed tracking-tight">
                {children}
            </li>
        ),
        blockquote: ({ children }) => (
            <blockquote className="border-l-4 border-border pl-4 italic text-muted my-6">
                {children}
            </blockquote>
        ),
        hr: () => <hr className="h-px w-full bg-border my-8 border-0" />,
        table: ({ children }) => (
            <div className="overflow-x-auto mb-6">
                <table className="w-full text-base border-collapse">
                    {children}
                </table>
            </div>
        ),
        thead: ({ children }) => (
            <thead>{children}</thead>
        ),
        tbody: ({ children }) => (
            <tbody>{children}</tbody>
        ),
        tr: ({ children }) => (
            <tr>{children}</tr>
        ),
        th: ({ children }) => (
            <th className="text-left font-medium text-base-content border-b border-border px-3 py-2">
                {children}
            </th>
        ),
        td: ({ children }) => (
            <td className="text-muted border-b border-border px-3 py-2">
                {children}
            </td>
        ),
        ...components,
    }
}
