"use client"

import {
  Children,
  cloneElement,
  createContext,
  isValidElement,
  useCallback,
  useContext,
  useId,
  useState,
  type ReactElement,
  type ReactNode,
} from "react"
import { MenubarSub, MenubarSubContent } from "@/components/ui/menubar"

type ExclusiveMenuSubContextValue = {
  openId: string | null
  setOpenId: (id: string | null) => void
}

const ExclusiveMenuSubContext = createContext<ExclusiveMenuSubContextValue | null>(null)

/** Ensures only one MenubarSub stays open while hovering across sibling sub-triggers. */
export function DocsExclusiveMenuSubRoot({ children }: { children: ReactNode }) {
  const [openId, setOpenId] = useState<string | null>(null)
  return (
    <ExclusiveMenuSubContext.Provider value={{ openId, setOpenId }}>
      {children}
    </ExclusiveMenuSubContext.Provider>
  )
}

function isMenubarSubContentElement(child: ReactNode): child is ReactElement<{ children?: ReactNode }> {
  return isValidElement(child) && child.type === MenubarSubContent
}

/** Nested exclusive subs get their own open-id scope so parents stay open. */
function withNestedExclusiveSubRoot(children: ReactNode): ReactNode {
  return Children.map(children, (child) => {
    if (!isMenubarSubContentElement(child)) return child

    return cloneElement(child, {
      children: <DocsExclusiveMenuSubRoot>{child.props.children}</DocsExclusiveMenuSubRoot>,
    })
  })
}

export function DocsExclusiveMenuSub({
  id,
  children,
}: {
  id: string
  children: ReactNode
}) {
  const ctx = useContext(ExclusiveMenuSubContext)
  const fallbackId = useId()
  const subId = id || fallbackId

  const open = ctx?.openId === subId
  const onOpenChange = useCallback(
    (next: boolean) => {
      if (!ctx) return
      if (next) {
        ctx.setOpenId(subId)
        return
      }
      // Only clear if this sub is still the active one (avoid race when switching siblings).
      if (ctx.openId === subId) {
        ctx.setOpenId(null)
      }
    },
    [ctx, subId]
  )

  if (!ctx) {
    return <MenubarSub>{children}</MenubarSub>
  }

  return (
    <MenubarSub open={open} onOpenChange={onOpenChange}>
      {withNestedExclusiveSubRoot(children)}
    </MenubarSub>
  )
}