import * as React from 'react'
import { ChevronRightIcon } from 'lucide-react'

import { cn } from './utils'

/* ------------------------------------------------------------------ */
/* Types                                                               */
/* ------------------------------------------------------------------ */

export interface TreeNode {
  id: string
  label: string
  icon?: React.ReactNode
  type?: string
  children?: TreeNode[]
  /** Whether the node can be expanded (defaults to true if children exist) */
  expandable?: boolean
  expanded?: boolean
  /** Async loader for lazy children */
  lazyLoad?: () => Promise<TreeNode[]>
  /** Arbitrary user data */
  data?: unknown
}

export interface TreeViewProps extends Omit<
  React.ComponentProps<'div'>,
  'onSelect' | 'onDoubleClick'
> {
  /** Tree data */
  data: TreeNode[]
  /** Callback when a node is selected */
  onSelect?: (node: TreeNode) => void
  /** Callback when a node is double-clicked */
  onDoubleClick?: (node: TreeNode) => void
  /** Callback when a node is expanded */
  onExpand?: (node: TreeNode) => void
  /** The id of the selected node (controlled) */
  selectedId?: string | null
  /** Default expanded node ids */
  defaultExpandedIds?: string[]
  /** Whether to show a search/filter input */
  searchable?: boolean
  /** Custom class for items */
  itemClassName?: string
}

/* ------------------------------------------------------------------ */
/* TreeView                                                            */
/* ------------------------------------------------------------------ */

function TreeView({
  className,
  data,
  onSelect,
  onDoubleClick,
  onExpand,
  selectedId: controlledSelectedId,
  defaultExpandedIds,
  searchable = false,
  itemClassName,
  ...props
}: TreeViewProps) {
  const [expanded, setExpanded] = React.useState<Set<string>>(
    () => new Set(defaultExpandedIds ?? [])
  )
  const [internalSelectedId, setInternalSelectedId] = React.useState<string | null>(null)
  const selectedId = controlledSelectedId !== undefined ? controlledSelectedId : internalSelectedId

  const [filterQuery, setFilterQuery] = React.useState('')
  const [loadingNodes, setLoadingNodes] = React.useState<Set<string>>(new Set())
  const [lazyChildren, setLazyChildren] = React.useState<Map<string, TreeNode[]>>(new Map())

  const toggleExpand = React.useCallback(
    (node: TreeNode) => {
      const isNodeExpanded = expanded.has(node.id)

      if (isNodeExpanded) {
        setExpanded(prev => {
          const next = new Set(prev)
          next.delete(node.id)
          return next
        })
      } else {
        // Lazy load children if needed (guard against duplicate in-flight requests)
        if (
          node.lazyLoad !== undefined &&
          (node.children === undefined || node.children.length === 0) &&
          !lazyChildren.has(node.id) &&
          !loadingNodes.has(node.id)
        ) {
          setLoadingNodes(prev => new Set(prev).add(node.id))
          void node
            .lazyLoad()
            .then(children => {
              setLazyChildren(prev => new Map(prev).set(node.id, children))
              setExpanded(prev => new Set(prev).add(node.id))
              onExpand?.(node)
            })
            .finally(() => {
              setLoadingNodes(prev => {
                const next = new Set(prev)
                next.delete(node.id)
                return next
              })
            })
        } else {
          setExpanded(prev => new Set(prev).add(node.id))
          onExpand?.(node)
        }
      }
    },
    [expanded, onExpand, lazyChildren]
  )

  const handleSelect = React.useCallback(
    (node: TreeNode) => {
      if (controlledSelectedId === undefined) setInternalSelectedId(node.id)
      onSelect?.(node)
    },
    [controlledSelectedId, onSelect]
  )

  // Filter matcher
  const matchesFilter = React.useCallback((node: TreeNode, q: string): boolean => {
    if (q === '') return true
    if (node.label.toLowerCase().includes(q)) return true
    if (node.children !== undefined) {
      return node.children.some(child => matchesFilter(child, q))
    }
    return false
  }, [])

  const query = filterQuery.toLowerCase()

  return (
    <div data-slot="tree-view" className={cn('flex flex-col gap-1', className)} {...props}>
      {searchable && (
        <div className="px-1 pb-1">
          <input
            type="text"
            placeholder="Filter..."
            value={filterQuery}
            onChange={e => {
              setFilterQuery(e.target.value)
            }}
            aria-label="Filter tree items"
            className="h-8 w-full rounded-lg border border-border-glass bg-glass px-3 text-xs text-foreground placeholder:text-foreground-muted backdrop-blur-[20px] outline-none focus:border-primary focus:ring-1 focus:ring-primary/20"
          />
        </div>
      )}

      <div role="tree" aria-label="Tree navigation" className="flex flex-col">
        {data.map(node =>
          matchesFilter(node, query) ? (
            <TreeItem
              key={node.id}
              node={node}
              depth={0}
              expanded={expanded}
              selectedId={selectedId}
              onToggle={toggleExpand}
              onSelect={handleSelect}
              onDoubleClick={onDoubleClick}
              matchesFilter={matchesFilter}
              filterQuery={query}
              loadingNodes={loadingNodes}
              lazyChildren={lazyChildren}
              itemClassName={itemClassName}
            />
          ) : null
        )}
      </div>
    </div>
  )
}

/* ------------------------------------------------------------------ */
/* TreeItem (recursive internal)                                       */
/* ------------------------------------------------------------------ */

interface TreeItemProps {
  node: TreeNode
  depth: number
  expanded: Set<string>
  selectedId: string | null
  onToggle: (node: TreeNode) => void
  onSelect: (node: TreeNode) => void
  onDoubleClick?: ((node: TreeNode) => void) | undefined
  matchesFilter: (node: TreeNode, query: string) => boolean
  filterQuery: string
  loadingNodes: Set<string>
  lazyChildren: Map<string, TreeNode[]>
  itemClassName?: string | undefined
}

function TreeItem({
  node,
  depth,
  expanded,
  selectedId,
  onToggle,
  onSelect,
  onDoubleClick,
  matchesFilter,
  filterQuery,
  loadingNodes,
  lazyChildren,
  itemClassName,
}: TreeItemProps) {
  const resolvedChildren = lazyChildren.get(node.id) ?? node.children
  const hasChildren =
    (resolvedChildren !== undefined && resolvedChildren.length > 0) || node.lazyLoad !== undefined
  const isExpandable = node.expandable !== false && hasChildren
  const isExpanded = expanded.has(node.id)
  const isSelected = selectedId === node.id
  const isLoading = loadingNodes.has(node.id)

  return (
    <div data-slot="tree-item" data-node-id={node.id}>
      <div
        role="treeitem"
        aria-expanded={isExpandable ? isExpanded : undefined}
        aria-selected={isSelected}
        aria-level={depth + 1}
        tabIndex={isSelected ? 0 : -1}
        className={cn(
          'flex items-center gap-1 rounded-lg px-2 py-1 text-sm cursor-pointer select-none transition-colors',
          'text-foreground hover:bg-glass-hover',
          isSelected && 'bg-primary/10 text-primary',
          itemClassName
        )}
        style={{ paddingLeft: `${String(depth * 16 + 8)}px` }}
        onClick={() => {
          onSelect(node)
          if (isExpandable) onToggle(node)
        }}
        onDoubleClick={() => {
          onDoubleClick?.(node)
        }}
        onKeyDown={e => {
          if (e.key === 'Enter' || e.key === ' ') {
            e.preventDefault()
            onSelect(node)
          }
          if (e.key === 'ArrowRight' && isExpandable && !isExpanded) {
            e.preventDefault()
            onToggle(node)
          }
          if (e.key === 'ArrowLeft' && isExpandable && isExpanded) {
            e.preventDefault()
            onToggle(node)
          }
        }}
      >
        {/* Expand/collapse chevron */}
        <span className="flex size-4 shrink-0 items-center justify-center">
          {isExpandable ? (
            isLoading ? (
              <span className="size-3 animate-spin rounded-full border-2 border-foreground-muted border-t-transparent" />
            ) : (
              <ChevronRightIcon
                className={cn(
                  'size-3.5 text-foreground-muted transition-transform duration-200',
                  isExpanded && 'rotate-90'
                )}
              />
            )
          ) : null}
        </span>

        {/* Icon */}
        {node.icon !== undefined && <span className="shrink-0">{node.icon}</span>}

        {/* Label */}
        <span className="truncate">{node.label}</span>
      </div>

      {/* Children */}
      {isExpanded && resolvedChildren !== undefined && resolvedChildren.length > 0 && (
        <div role="group">
          {resolvedChildren.map(child =>
            matchesFilter(child, filterQuery) ? (
              <TreeItem
                key={child.id}
                node={child}
                depth={depth + 1}
                expanded={expanded}
                selectedId={selectedId}
                onToggle={onToggle}
                onSelect={onSelect}
                onDoubleClick={onDoubleClick}
                matchesFilter={matchesFilter}
                filterQuery={filterQuery}
                loadingNodes={loadingNodes}
                lazyChildren={lazyChildren}
                itemClassName={itemClassName}
              />
            ) : null
          )}
        </div>
      )}
    </div>
  )
}

export { TreeView, TreeItem }
