"use client";

import {
	AnimatePresence,
	HTMLMotionProps,
	type Variants,
	backIn,
	backOut,
	motion,
} from "motion/react";
import {
	Children,
	cloneElement,
	createContext,
	type Dispatch,
	isValidElement,
	type KeyboardEvent,
	type ReactElement,
	type ReactNode,
	type RefObject,
	type SetStateAction,
	useCallback,
	useContext,
	useEffect,
	useMemo,
	useRef,
	useState,
} from "react";
import { cn } from "@/lib/utils";

/* ------------------------- Easings + variants ----------------------------- */

const expoIn = (t: number) => (t === 0 ? 0 : 2 ** (10 * (t - 1)));
const expoOut = (t: number) => (t === 1 ? 1 : 1 - 2 ** (-10 * t));

const navItemVariants: Variants = {
	initial: { opacity: 0, filter: "blur(10px)", translateX: 40 },
	open: ({ index }: { index: number }) => ({
		opacity: 1,
		filter: "blur(0px)",
		translateX: 0,
		transition: { duration: 0.6, ease: expoOut, delay: 0.3 + index * 0.1 },
	}),
	closed: ({
		index,
		secondaryPanelOpen,
	}: { index: number; secondaryPanelOpen: boolean }) => ({
		opacity: 0,
		filter: "blur(10px)",
		translateX: 40,
		transition: {
			duration: 0.6,
			ease: (t: number) => expoIn(t),
			delay: 0.05 + (secondaryPanelOpen ? 0.15 : 0) + index * 0.1,
		},
	}),
};

const mainPanelVariants: Variants = {
	initial: { opacity: 0, filter: "blur(10px)", translateY: 20 },
	animate: {
		opacity: 1,
		filter: "blur(0px)",
		translateY: 0,
		transition: { duration: 0.5, ease: expoOut },
	},
	exit: ({ isSecondaryPanelOpen }: { isSecondaryPanelOpen: boolean }) => ({
		opacity: 0,
		filter: "blur(10px)",
		translateY: 10,
		transition: {
			duration: 0.5,
			ease: expoIn,
			delay: isSecondaryPanelOpen ? 0.1 : 0,
		},
	}),
};

const secondaryPanelVariants: Variants = {
	initial: { opacity: 0, filter: "blur(4px)", translateY: 5 },
	animate: {
		opacity: 1,
		filter: "blur(0px)",
		translateY: 0,
		transition: { duration: 0.6, type: "spring", bounce: 0 },
	},
	exit: (isMenuOpen: boolean) => ({
		opacity: 0,
		filter: isMenuOpen ? "blur(4px)" : "blur(10px)",
		translateY: isMenuOpen ? 5 : 10,
		transition: isMenuOpen
			? { duration: 0.6, type: "spring", bounce: 0 }
			: { duration: 0.5, ease: expoIn },
	}),
};

/* --------------------------------- Context -------------------------------- */

type FloatingMenuContextValue = {
	isMenuOpen: boolean;
	toggleMenu: () => void;
	closeMenu: () => void;
	activeIndex: number | null;
	setActiveIndex: Dispatch<SetStateAction<number | null>>;
	focusedIndex: number | null;
	setFocusedIndex: Dispatch<SetStateAction<number | null>>;
	forcedFocusVisibleIndex: number | null;
	setForcedFocusVisibleIndex: Dispatch<SetStateAction<number | null>>;
	triggerRef: RefObject<HTMLButtonElement | null>;
	itemsRefs: RefObject<HTMLButtonElement[]>;
	openedWithPointerRef: RefObject<boolean>;
};

const FloatingMenuContext = createContext<FloatingMenuContextValue | null>(null);

const useFloatingMenu = () => {
	const ctx = useContext(FloatingMenuContext);
	if (!ctx) {
		throw new Error("FloatingMenu.* must be used inside <FloatingMenu.Root>");
	}
	return ctx;
};

/* --------------------------------- Trigger -------------------------------- */

const Trigger = ({ className, ...props }: HTMLMotionProps<"button">) => {
	const { isMenuOpen, toggleMenu, triggerRef, openedWithPointerRef } =
		useFloatingMenu();

	const handlePointerDown = () => {
		openedWithPointerRef.current = true;
	};

	return (
		<motion.button
			ref={triggerRef}
			type="button"
			whileTap={{ scale: 0.95 }}
			onPointerDown={handlePointerDown}
			onClick={toggleMenu}
			aria-expanded={isMenuOpen}
			aria-haspopup="true"
            data-focused={document && (triggerRef.current === document.activeElement)}
            data-open={isMenuOpen}
			aria-controls="floating-menu-main"
			aria-label={isMenuOpen ? "Close menu" : "Open menu"}
			className={cn(
				"size-8 md:size-10 text-sm cursor-pointer relative z-10 aspect-square rounded-full backdrop-blur-sm pointer-events-auto grid place-items-center",
				"bg-black/10 dark:bg-white/5 border border-black/20 dark:border-white/5 text-[#2b2a2a] dark:text-white hover:border-black/50 dark:hover:border-white/20",
				"transition-colors duration-500",
				"data-[focused=true]:outline-none data-[focused=true]:border-black data-[focused=true]:dark:border-white/20",
				className
			)}
            {...props}
		>
			<AnimatePresence mode="wait" initial={false}>
				{isMenuOpen ? (
					<motion.svg
						key="close-icon"
						viewBox="0 0 24 24"
						className="size-3 [grid-area:1/1]"
						initial={{ rotate: -45 }}
						animate={{ rotate: 0 }}
						exit={{ rotate: -45 }}
						transition={{ duration: 0.5, ease: backOut }}
					>
						<motion.path
							strokeWidth="3"
							d="M 0,0 L 24,24"
							stroke="currentColor"
						/>
						<motion.path
							strokeWidth="3"
							d="M 24,0 L 0,24"
							stroke="currentColor"
							initial={{ rotate: 90 }}
							animate={{ rotate: 0 }}
							exit={{ rotate: 90 }}
							transition={{ duration: 0.25, ease: backOut }}
						/>
					</motion.svg>
				) : (
					<motion.svg
						key="open-icon"
						viewBox="0 0 24 24"
						className="size-4 [grid-area:1/1]"
						transition={{ duration: 0.5, ease: backIn }}
					>
						<motion.path
							strokeWidth="3"
							d="M 0,3 L 24,3"
							stroke="currentColor"
							strokeLinecap="round"
							initial={{ opacity: 0, translateY: 10 }}
							animate={{ opacity: 1, translateY: 0 }}
							exit={{ opacity: 0, translateY: 10 }}
						/>
						<motion.path
							strokeWidth="3"
							d="M 0,12 L 24,12"
							stroke="currentColor"
							strokeLinecap="round"
						/>
						<motion.path
							strokeWidth="3"
							d="M 0,21 L 24,21"
							stroke="currentColor"
							strokeLinecap="round"
							initial={{ opacity: 0, translateY: -10 }}
							animate={{ opacity: 1, translateY: 0 }}
							exit={{ opacity: 0, translateY: -10 }}
						/>
					</motion.svg>
				)}
			</AnimatePresence>
		</motion.button>
	);
};
Trigger.displayName = "FloatingMenu.Trigger";

/* ----------------------------------- Nav ---------------------------------- */

type NavProps = {
	children: ReactNode;
	className?: string;
};

const Nav = ({ children, className }: NavProps) => {
	const { isMenuOpen, itemsRefs, setForcedFocusVisibleIndex, triggerRef } =
		useFloatingMenu();

	const childArray = Children.toArray(children).filter(isValidElement);

	const handleKeyDown = (e: KeyboardEvent) => {
		setForcedFocusVisibleIndex(null);

		if (e.key === "Tab") {
			e.preventDefault();
			triggerRef.current?.focus();
			return;
		}

		if (e.key === "ArrowLeft" || e.key === "ArrowRight") {
			e.preventDefault();
			const currentIndex = itemsRefs.current.findIndex(
				(r) => r === document.activeElement
			);
			if (currentIndex === -1 || childArray.length === 0) {
				return;
			}
			const direction = e.key === "ArrowRight" ? 1 : -1;
			const nextIndex =
				(currentIndex + direction + childArray.length) % childArray.length;
			itemsRefs.current[nextIndex]?.focus();
		}
	};

	return (
		<AnimatePresence>
			{isMenuOpen && (
				<motion.nav
					aria-label="Menu options"
					onKeyDown={handleKeyDown}
					className={cn("flex items-center gap-2 flex-1", className)}
				>
					{childArray.map((child, index) =>
						cloneElement(child as ReactElement<any>, { index, key: index })
					)}
				</motion.nav>
			)}
		</AnimatePresence>
	);
};
Nav.displayName = "FloatingMenu.Nav";

/* --------------------------------- NavItem -------------------------------- */

type NavItemProps = {
	children?: ReactNode;
	className?: string;
	/** Injected by Nav, do not pass manually */
	index?: number;
};

const NavItem = ({ children, className, index = 0 }: NavItemProps) => {
	const {
		activeIndex,
		setActiveIndex,
		itemsRefs,
		focusedIndex,
		setFocusedIndex,
		forcedFocusVisibleIndex,
		setForcedFocusVisibleIndex,
	} = useFloatingMenu();

	const isActive = activeIndex === index;

	const setRef = useCallback(
		(el: HTMLButtonElement | null) => {
			if (el) {
				itemsRefs.current[index] = el;
			}
		},
		[index, itemsRefs]
	);

	const handleClick = () => {
		if (!isActive) {
			setActiveIndex(index);
		}
	};

	return (
		<motion.button
			ref={setRef}
			type="button"
			variants={navItemVariants}
			initial="closed"
			animate="open"
			exit="closed"
			custom={{ index, secondaryPanelOpen: activeIndex !== null }}
			onClick={handleClick}
			onPointerDown={() => setForcedFocusVisibleIndex(null)}
			onFocus={() => setFocusedIndex(index)}
			onBlur={() =>
				setFocusedIndex((prev) => (prev === index ? null : prev))
			}
			aria-expanded={isActive}
			aria-controls={`floating-menu-secondary-${index}`}
			aria-haspopup="true"
			aria-selected={isActive}

            data-focused={forcedFocusVisibleIndex === index && focusedIndex === index}
            data-active={isActive}
			className={cn(
				"rounded-xl flex-1 py-1",
				"cursor-pointer transition-colors duration-500",
				"text-sm dark:text-white text-[#2b2a2a] z-5 backdrop-blur-sm",
				"bg-black/10 dark:bg-white/5",
				"border border-black/20 dark:border-white/5 hover:border-black/50 dark:hover:border-white/20",
				"focus:outline-none focus-visible:border-black focus-visible:dark:border-white/20",
				"data-[focused=true]:border-black data-[focused=true]:dark:border-white/20",
				"data-[active=true]:dark:bg-white/10 data-[active=true]:bg-black/20",
				className
			)}
		>
			{children}
		</motion.button>
	);
};
NavItem.displayName = "FloatingMenu.NavItem";

/* -------------------------------- MainPanel ------------------------------- */

type MainPanelProps = {
	children?: ReactNode;
	className?: string;
};

const MainPanel = ({ children, className }: MainPanelProps) => {
	const { isMenuOpen, activeIndex } = useFloatingMenu();

	return (
		<AnimatePresence>
			{isMenuOpen && (
				<motion.div
					key="main-panel"
					id="floating-menu-main"
					role="region"
					aria-label="Core panel"
					custom={{ isSecondaryPanelOpen: activeIndex !== null }}
					variants={mainPanelVariants}
					initial="initial"
					animate="animate"
					exit="exit"
					className={cn(
						"w-full rounded-xl overflow-hidden backdrop-blur-sm",
						"bg-black/10 dark:bg-white/5 border border-black/20 dark:border-white/5",
						"text-sm text-[#2b2a2a] dark:text-white flex flex-col gap-2 p-3",
						"transition-colors duration-500",
						className
					)}
				>
					{children}
				</motion.div>
			)}
		</AnimatePresence>
	);
};
MainPanel.displayName = "FloatingMenu.MainPanel";

/* ----------------------------- SecondaryPanel ----------------------------- */

type SecondaryPanelProps = {
	children?: ReactNode;
	className?: string;
};

const SecondaryPanel = ({ children, className }: SecondaryPanelProps) => {
	const { activeIndex, isMenuOpen, setActiveIndex } = useFloatingMenu();

	const items = Children.toArray(children).filter(isValidElement);
	const defaultOpenItemIndex = items.findIndex(
		(child) => (child.props as SecondaryPanelItemProps).defaultOpen
	);

	useEffect(() => {
		if (isMenuOpen && defaultOpenItemIndex > -1) {
			setActiveIndex(defaultOpenItemIndex);
		}
	}, [isMenuOpen, defaultOpenItemIndex, setActiveIndex]);

	const activeItem =
		activeIndex !== null && activeIndex < items.length
			? items[activeIndex]
			: null;

	return (
		<AnimatePresence mode="wait" custom={isMenuOpen}>
			{activeItem && isMenuOpen && (
				<motion.div
                    role="region"
                    aria-label={`Secondary panel ${activeIndex !== null && `for option ${activeIndex + 1}` }`}
					key={activeIndex}
					id={`floating-menu-secondary-${activeIndex}`}
					aria-live="polite"
					custom={isMenuOpen}
					variants={secondaryPanelVariants}
					initial="initial"
					animate="animate"
					exit="exit"
					className={cn(
						"text-sm rounded-xl p-3 bg-black/10 dark:bg-white/5 backdrop-blur-sm",
						"border border-black/20 dark:border-white/5 w-full",
						"text-[#2b2a2a] dark:text-white",
						className
					)}
				>
					{activeItem}
				</motion.div>
			)}
		</AnimatePresence>
	);
};
SecondaryPanel.displayName = "FloatingMenu.SecondaryPanel";

/* --------------------------- SecondaryPanelItem --------------------------- */

type SecondaryPanelItemProps = {
	children?: ReactNode;
	className?: string;
	/** When true, this item is auto-activated whenever the menu opens. */
	defaultOpen?: boolean;
};

const SecondaryPanelItem = ({ children, className }: SecondaryPanelItemProps) => {
	if (className) {
		return <div className={className}>{children}</div>;
	}
	return children;
};
SecondaryPanelItem.displayName = "FloatingMenu.SecondaryPanelItem";

/* ---------------------------------- Root ---------------------------------- */

const Root = ({ children, className, ...props }: React.ComponentProps<"div">) => {
	const [isMenuOpen, setIsMenuOpen] = useState(false);
	const [activeIndex, setActiveIndex] = useState<number | null>(null);
	const [focusedIndex, setFocusedIndex] = useState<number | null>(null);
	const [forcedFocusVisibleIndex, setForcedFocusVisibleIndex] = useState<
		number | null
	>(null);

	const triggerRef = useRef<HTMLButtonElement | null>(null);
	const itemsRefs = useRef<HTMLButtonElement[]>([]);
	const openedWithPointerRef = useRef(false);

	const toggleMenu = useCallback(() => {
		setIsMenuOpen((o) => !o);
		setActiveIndex(null);
	}, []);

	const closeMenu = useCallback(() => {
		setIsMenuOpen(false);
		setActiveIndex(null);
		triggerRef.current?.focus();
	}, []);

	useEffect(() => {
		if (isMenuOpen) {
			itemsRefs.current[0]?.focus();
			if (openedWithPointerRef.current) {
				setForcedFocusVisibleIndex(0);
			}
		} else {
			setFocusedIndex(null);
			setForcedFocusVisibleIndex(null);
		}
		openedWithPointerRef.current = false;
	}, [isMenuOpen]);

	let trigger: ReactNode = null;
	let nav: ReactNode = null;
	let main: ReactNode = null;
	let secondary: ReactNode = null;
	Children.forEach(children, (child) => {
		if (!isValidElement(child)) {
			return;
		}
		const t = child.type;
		if (t === Trigger) {
			trigger = child;
		} else if (t === Nav) {
			nav = child;
		} else if (t === MainPanel) {
			main = child;
		} else if (t === SecondaryPanel) {
			secondary = child;
		}
	});

	const handleKeyDown = (e: KeyboardEvent) => {
		if (e.key === "Escape") {
			closeMenu();
		}
	};

	const value = useMemo<FloatingMenuContextValue>(
		() => ({
			isMenuOpen,
			toggleMenu,
			closeMenu,
			activeIndex,
			setActiveIndex,
			focusedIndex,
			setFocusedIndex,
			forcedFocusVisibleIndex,
			setForcedFocusVisibleIndex,
			triggerRef,
			itemsRefs,
			openedWithPointerRef,
		}),
		[
			isMenuOpen,
			toggleMenu,
			closeMenu,
			activeIndex,
			focusedIndex,
			forcedFocusVisibleIndex,
		]
	);

	return (
		<FloatingMenuContext.Provider value={value}>
			<div
				role="menu"
				aria-label="Main Menu"
				className={cn(
					"absolute z-99",
					"flex flex-col items-end gap-2 w-84",
					!isMenuOpen && "pointer-events-none",
					className
				)}
                {...props}
				onKeyDown={handleKeyDown}
			>
				<div className="flex items-center justify-end gap-2 w-full">
					{nav}
					{trigger}
				</div>
				{main}
				{secondary}
			</div>
		</FloatingMenuContext.Provider>
	);
};
Root.displayName = "FloatingMenu.Root";

/* --------------------------------- Export --------------------------------- */

const FloatingMenu = {
	Root,
	Trigger,
	Nav,
	NavItem,
	MainPanel,
	SecondaryPanel,
	SecondaryPanelItem,
};

export { FloatingMenu };

import { FloatingMenu } from "@/components/floating-menu";

return(
  <FloatingMenu.Root className="top-12 right-12">
    <FloatingMenu.Trigger />
    <FloatingMenu.Nav>
      {OPTIONS.map((opt) => (
        <FloatingMenu.NavItem key={opt.id}>{opt.label}</FloatingMenu.NavItem>
      ))}
    </FloatingMenu.Nav>
    <FloatingMenu.MainPanel>
      Main content panel
    </FloatingMenu.MainPanel>
    <FloatingMenu.SecondaryPanel>
      {OPTIONS.map((opt) => (
        <FloatingMenu.SecondaryPanelItem key={opt.id}>
          {opt.content}
        </FloatingMenu.SecondaryPanelItem>
      ))}
    </FloatingMenu.SecondaryPanel>
  </FloatingMenu.Root>
)

<FloatingMenu.Root className="top-12 right-12">
  <FloatingMenu.Trigger />
  <FloatingMenu.Nav>
    {OPTIONS.map((opt) => (
      <FloatingMenu.NavItem key={opt.id}>{opt.label}</FloatingMenu.NavItem>
    ))}
  </FloatingMenu.Nav>
  <FloatingMenu.MainPanel>
    Main content panel
  </FloatingMenu.MainPanel>
  <FloatingMenu.SecondaryPanel>
    {OPTIONS.map((opt) => (
      <FloatingMenu.SecondaryPanelItem key={opt.id} defaultOpen={opt.id === "o1"}>
        {opt.content}
      </FloatingMenu.SecondaryPanelItem>
    ))}
  </FloatingMenu.SecondaryPanel>
</FloatingMenu.Root>

<FloatingMenu.Root className="top-12 right-12">
  <FloatingMenu.Trigger />
  <FloatingMenu.Nav>
    {OPTIONS.map((opt) => (
      <FloatingMenu.NavItem key={opt.id}>{opt.label}</FloatingMenu.NavItem>
    ))}
  </FloatingMenu.Nav>
  {showMainPanel && (
    <FloatingMenu.MainPanel>
      Main content panel
    </FloatingMenu.MainPanel>
  )}
  <FloatingMenu.SecondaryPanel>
    {OPTIONS.map((opt) => (
      <FloatingMenu.SecondaryPanelItem key={opt.id}>
        {opt.content}
      </FloatingMenu.SecondaryPanelItem>
    ))}
  </FloatingMenu.SecondaryPanel>
</FloatingMenu.Root>