import { Fragment, useCallback, useEffect, useRef, useState } from 'react'
import { motion, AnimatePresence } from 'framer-motion'
import { ChevronDownIcon, ChevronRightIcon } from '@heroicons/react/24/outline'
import { twMerge } from 'tailwind-merge'

export interface TreeNode {
    id: string | number
    label: string
    children?: TreeNode[]
    disabled?: boolean
}

export interface TreeSelectProps {
    data: TreeNode[]
    value?: string | number
    onChange?: (value: string | number) => void
    placeholder?: string
    disabled?: boolean
    className?: string
    dropdownClassName?: string
    notFoundContent?: React.ReactNode
    loading?: boolean
    error?: string
}

export default function TreeSelect({
    data,
    value,
    onChange,
    placeholder = 'Please select...',
    disabled = false,
    className,
    dropdownClassName,
    notFoundContent = 'No data',
    loading = false,
    error,
}: TreeSelectProps) {
    const [isOpen, setIsOpen] = useState(false)
    const [expandedKeys, setExpandedKeys] = useState<Set<string | number>>(new Set())
    const containerRef = useRef<HTMLDivElement>(null)
    const [selectedLabel, setSelectedLabel] = useState<string>('')

    // Close dropdown when clicking outside
    useEffect(() => {
        const handleClickOutside = (event: MouseEvent) => {
            if (containerRef.current && !containerRef.current.contains(event.target as Node)) {
                setIsOpen(false)
            }
        }

        document.addEventListener('mousedown', handleClickOutside)
        return () => document.removeEventListener('mousedown', handleClickOutside)
    }, [])

    // Find and set selected label when value changes
    useEffect(() => {
        const findLabel = (nodes: TreeNode[]): string => {
            for (const node of nodes) {
                if (node.id === value) return node.label
                if (node.children) {
                    const childLabel = findLabel(node.children)
                    if (childLabel) return childLabel
                }
            }
            return ''
        }
        setSelectedLabel(findLabel(data))
    }, [value, data])

    const toggleExpand = useCallback((key: string | number) => {
        setExpandedKeys(prev => {
            const next = new Set(prev)
            if (next.has(key)) {
                next.delete(key)
            } else {
                next.add(key)
            }
            return next
        })
    }, [])

    const handleSelect = useCallback((node: TreeNode) => {
        if (node.disabled) return
        onChange?.(node.id)
        setIsOpen(false)
    }, [onChange])

    const renderNode = useCallback((node: TreeNode, level = 0) => {
        const hasChildren = node.children && node.children.length > 0
        const isExpanded = expandedKeys.has(node.id)

        return (
            <Fragment key={node.id}>
                <motion.div
                    className={twMerge(
                        'flex items-center px-4 py-2 cursor-pointer hover:bg-gray-100 dark:hover:bg-gray-700',
                        'transition-colors duration-150 ease-out',
                        node.disabled && 'opacity-50 cursor-not-allowed',
                        value === node.id && 'bg-blue-50 dark:bg-blue-900',
                        level > 0 && 'pl-[calc(1rem*var(--level))]'
                    )}
                    style={{ '--level': level + 1 } as React.CSSProperties}
                    onClick={() => !hasChildren && handleSelect(node)}
                    role="option"
                    aria-selected={value === node.id}
                    tabIndex={0}
                    initial={{ opacity: 0, y: -10 }}
                    animate={{ opacity: 1, y: 0 }}
                    exit={{ opacity: 0, y: -10 }}
                    transition={{ duration: 0.2 }}
                >
                    {hasChildren && (
                        <motion.button
                            onClick={(e) => {
                                e.stopPropagation()
                                toggleExpand(node.id)
                            }}
                            className="p-1 hover:bg-gray-200 dark:hover:bg-gray-600 rounded-full mr-2"
                            aria-expanded={isExpanded}
                            whileHover={{ scale: 1.1 }}
                            whileTap={{ scale: 0.95 }}
                        >
                            <motion.div
                                initial={false}
                                animate={{ rotate: isExpanded ? 90 : 0 }}
                                transition={{ duration: 0.2 }}
                            >
                                <ChevronRightIcon className="h-4 w-4" />
                            </motion.div>
                        </motion.button>
                    )}
                    {!hasChildren && <div className="w-6 mr-2" />}
                    <span className="truncate">{node.label}</span>
                </motion.div>
                <AnimatePresence>
                    {hasChildren && isExpanded && (
                        <motion.div
                            role="group"
                            initial={{ opacity: 0, height: 0 }}
                            animate={{ opacity: 1, height: 'auto' }}
                            exit={{ opacity: 0, height: 0 }}
                            transition={{ duration: 0.2 }}
                        >
                            {node.children.map(child => renderNode(child, level + 1))}
                        </motion.div>
                    )}
                </AnimatePresence>
            </Fragment>
        )
    }, [expandedKeys, handleSelect, toggleExpand, value])

    return (
        <div ref={containerRef} className={twMerge('relative', className)}>
            <div
                className={twMerge(
                    'w-full rounded-lg border bg-white px-4 py-2',
                    'transition-all duration-200 ease-out',
                    'flex items-center justify-between cursor-pointer select-none',
                    error
                        ? 'border-red-500 focus:border-red-500 focus:ring-2 focus:ring-red-200'
                        : 'border-gray-200 hover:border-gray-300 focus:border-blue-500 focus:ring-2 focus:ring-blue-100',
                    disabled && 'opacity-50 cursor-not-allowed',
                    isOpen && !error && 'border-blue-500 ring-2 ring-blue-100'
                )}
                onClick={() => !disabled && setIsOpen(!isOpen)}
                role="combobox"
                aria-expanded={isOpen}
                aria-haspopup="listbox"
                tabIndex={disabled ? -1 : 0}
            >
                <span className={twMerge('truncate', !selectedLabel && 'text-tertiary-400')}>
                    {selectedLabel || placeholder}
                </span>
                <motion.div
                    animate={{ rotate: isOpen ? 180 : 0 }}
                    transition={{ duration: 0.2 }}
                >
                    <ChevronDownIcon className="h-5 w-5 text-tertiary-400" />
                </motion.div>
            </div>

            <AnimatePresence>
                {isOpen && (
                    <motion.div
                        className={twMerge(
                            'absolute z-50 w-full mt-1 bg-white dark:bg-gray-800',
                            'border border-gray-200 dark:border-gray-600 rounded-lg shadow-lg',
                            'max-h-60 overflow-auto',
                            dropdownClassName
                        )}
                        role="listbox"
                        initial={{ opacity: 0, y: -10 }}
                        animate={{ opacity: 1, y: 0 }}
                        exit={{ opacity: 0, y: -10 }}
                        transition={{ duration: 0.2 }}
                    >
                        {loading ? (
                            <div className="flex items-center justify-center py-4">
                                <motion.div
                                    className="w-6 h-6 border-2 border-blue-500 border-t-transparent rounded-full"
                                    animate={{ rotate: 360 }}
                                    transition={{ duration: 1, repeat: Infinity, ease: "linear" }}
                                />
                            </div>
                        ) : data.length > 0 ? (
                            data.map(node => renderNode(node))
                        ) : (
                            <div className="px-4 py-2 text-tertiary-400 text-center">{notFoundContent}</div>
                        )}
                    </motion.div>
                )}
            </AnimatePresence>
        </div>
    )
}
