Chain of Thought

"use client";

import {
    createContext,
    useContext,
    useId,
    useState,
    isValidElement,
    Children,
    type ComponentPropsWithoutRef,
    type ReactElement,
    type ReactNode,
} from "react";
import { AnimatePresence, motion } from "framer-motion";
import {
    BrainCircuit,
    Check,
    ChevronDown,
    ImageIcon,
    Link2,
    LoaderCircle,
    type LucideIcon,
} from "lucide-react";
import { cn } from "@/lib/utils";

export type ChainOfThoughtStatus = "complete" | "active" | "pending";

type ChainOfThoughtContextValue = {
    open: boolean;
    setOpen: (open: boolean) => void;
    contentId: string;
    stepCount: number;
    completeCount: number;
    activeCount: number;
};

const ChainOfThoughtContext = createContext<ChainOfThoughtContextValue | null>(null);

function useChainOfThoughtContext(componentName: string) {
    const context = useContext(ChainOfThoughtContext);

    if (!context) {
        throw new Error(`${componentName} must be used inside <ChainOfThought />.`);
    }

    return context;
}

function collectStepStats(children: ReactNode): {
    stepCount: number;
    completeCount: number;
    activeCount: number;
} {
    let stepCount = 0;
    let completeCount = 0;
    let activeCount = 0;

    Children.forEach(children, (child) => {
        if (!isValidElement(child)) {
            return;
        }

        const element = child as ReactElement<{
            status?: ChainOfThoughtStatus;
            children?: ReactNode;
        }>;
        const elementType = element.type as { displayName?: string };
        const status = element.props.status;

        if (elementType.displayName === "ChainOfThoughtStep") {
            stepCount += 1;

            if (status === "complete") {
                completeCount += 1;
            }

            if (status === "active") {
                activeCount += 1;
            }
        }

        if (element.props.children) {
            const nested = collectStepStats(element.props.children);
            stepCount += nested.stepCount;
            completeCount += nested.completeCount;
            activeCount += nested.activeCount;
        }
    });

    return { stepCount, completeCount, activeCount };
}

function getStatusClasses(status: ChainOfThoughtStatus) {
    switch (status) {
        case "complete":
            return {
                dot: "border-emerald-400/40 bg-emerald-500/15 text-emerald-300",
                badge: "border-emerald-500/20 bg-emerald-500/10 text-emerald-300",
                line: "bg-emerald-400/20",
                text: "text-zinc-100",
            };
        case "active":
            return {
                dot: "border-indigo-400/40 bg-indigo-500/15 text-indigo-200 shadow-[0_0_20px_rgba(99,102,241,0.22)]",
                badge: "border-indigo-500/20 bg-indigo-500/10 text-indigo-200",
                line: "bg-indigo-400/20",
                text: "text-white",
            };
        default:
            return {
                dot: "border-white/10 bg-white/[0.04] text-zinc-500",
                badge: "border-white/10 bg-white/[0.04] text-zinc-500",
                line: "bg-white/8",
                text: "text-zinc-300",
            };
    }
}

function getStatusLabel(status: ChainOfThoughtStatus) {
    switch (status) {
        case "complete":
            return "Complete";
        case "active":
            return "Active";
        default:
            return "Pending";
    }
}

function getStatusIcon(status: ChainOfThoughtStatus) {
    switch (status) {
        case "complete":
            return <Check className="size-3.5" />;
        case "active":
            return <LoaderCircle className="size-3.5 animate-[spin_2.6s_linear_infinite]" />;
        default:
            return <div className="size-1.5 rounded-full bg-current" />;
    }
}

export interface ChainOfThoughtProps extends ComponentPropsWithoutRef<"div"> {
    open?: boolean;
    defaultOpen?: boolean;
    onOpenChange?: (open: boolean) => void;
}

export function ChainOfThought({
    open: openProp,
    defaultOpen = false,
    onOpenChange,
    className,
    children,
    ...props
}: ChainOfThoughtProps) {
    const [internalOpen, setInternalOpen] = useState(defaultOpen);
    const contentId = useId();
    const open = openProp ?? internalOpen;
    const stats = collectStepStats(children);

    function setOpen(nextOpen: boolean) {
        if (openProp === undefined) {
            setInternalOpen(nextOpen);
        }

        onOpenChange?.(nextOpen);
    }

    return (
        <ChainOfThoughtContext.Provider
            value={{
                open,
                setOpen,
                contentId,
                stepCount: stats.stepCount,
                completeCount: stats.completeCount,
                activeCount: stats.activeCount,
            }}
        >
            <div
                data-state={open ? "open" : "closed"}
                className={cn(
                    "w-full max-w-3xl overflow-hidden rounded-[28px] border border-white/10 bg-[radial-gradient(circle_at_top,rgba(99,102,241,0.16),transparent_32%),linear-gradient(180deg,rgba(15,23,42,0.96),rgba(2,6,23,0.98))] p-3 text-zinc-100 shadow-[0_24px_80px_rgba(0,0,0,0.45)] backdrop-blur-xl",
                    className
                )}
                {...props}
            >
                {children}
            </div>
        </ChainOfThoughtContext.Provider>
    );
}

export interface ChainOfThoughtHeaderProps
    extends ComponentPropsWithoutRef<"button"> {
    children?: ReactNode;
}

export function ChainOfThoughtHeader({
    children,
    className,
    ...props
}: ChainOfThoughtHeaderProps) {
    const { open, setOpen, contentId, stepCount, completeCount, activeCount } =
        useChainOfThoughtContext("ChainOfThoughtHeader");

    return (
        <button
            type="button"
            aria-expanded={open}
            aria-controls={contentId}
            onClick={() => setOpen(!open)}
            className={cn(
                "flex w-full items-center justify-between gap-4 rounded-[22px] border border-white/8 bg-white/[0.03] px-4 py-3 text-left transition-colors hover:border-white/12 hover:bg-white/[0.05]",
                className
            )}
            {...props}
        >
            {children ?? (
                <div className="min-w-0">
                    <div className="mb-1 flex items-center gap-2">
                        <div className="flex size-9 items-center justify-center rounded-2xl border border-indigo-400/25 bg-indigo-500/12 text-indigo-200">
                            <BrainCircuit className="size-4.5" />
                        </div>
                        <div>
                            <div className="text-sm font-semibold text-white">
                                Chain of Thought
                            </div>
                            <div className="text-xs text-zinc-400">
                                {stepCount} step{stepCount === 1 ? "" : "s"} •{" "}
                                {completeCount} complete • {activeCount} active
                            </div>
                        </div>
                    </div>
                </div>
            )}

            <div className="flex items-center gap-2">
                <span className="hidden rounded-full border border-white/8 bg-white/[0.04] px-2.5 py-1 text-[11px] font-medium text-zinc-400 sm:inline-flex">
                    {open ? "Hide reasoning" : "Show reasoning"}
                </span>
                <motion.span
                    animate={{ rotate: open ? 180 : 0 }}
                    transition={{ duration: 0.22 }}
                    className="flex size-9 items-center justify-center rounded-2xl border border-white/8 bg-black/20 text-zinc-400"
                >
                    <ChevronDown className="size-4" />
                </motion.span>
            </div>
        </button>
    );
}

export type ChainOfThoughtContentProps = ComponentPropsWithoutRef<"div">;

export function ChainOfThoughtContent({
    className,
    children,
    ...props
}: ChainOfThoughtContentProps) {
    const { open, contentId } = useChainOfThoughtContext("ChainOfThoughtContent");

    return (
        <AnimatePresence initial={false}>
            {open && (
                <motion.div
                    id={contentId}
                    key="chain-of-thought-content"
                    initial={{ height: 0, opacity: 0 }}
                    animate={{ height: "auto", opacity: 1 }}
                    exit={{ height: 0, opacity: 0 }}
                    transition={{ duration: 0.28, ease: "easeOut" }}
                    className="overflow-hidden"
                >
                    <div
                        className={cn("px-2 pt-3 pb-1", className)}
                        {...props}
                    >
                        {children}
                    </div>
                </motion.div>
            )}
        </AnimatePresence>
    );
}

export interface ChainOfThoughtStepProps
    extends ComponentPropsWithoutRef<"div"> {
    icon?: LucideIcon;
    label?: string;
    description?: string;
    status?: ChainOfThoughtStatus;
}

export function ChainOfThoughtStep({
    icon: Icon = BrainCircuit,
    label,
    description,
    status = "pending",
    className,
    children,
    ...props
}: ChainOfThoughtStepProps) {
    const statusClasses = getStatusClasses(status);

    return (
        <div
            data-status={status}
            className={cn("relative pl-14 last:pb-0", className)}
            {...props}
        >
            <div
                className={cn(
                    "absolute left-[18px] top-12 bottom-0 w-px",
                    statusClasses.line
                )}
            />

            <div
                className={cn(
                    "absolute left-0 top-1 flex size-10 items-center justify-center rounded-2xl border",
                    statusClasses.dot
                )}
            >
                <Icon className="size-4" />
            </div>

            <div className="rounded-[22px] border border-white/8 bg-white/[0.03] px-4 py-3.5">
                <div className="mb-2 flex items-start justify-between gap-3">
                    <div className="min-w-0">
                        {label && (
                            <div
                                className={cn(
                                    "text-sm leading-6",
                                    statusClasses.text
                                )}
                            >
                                {label}
                            </div>
                        )}
                        {description && (
                            <p className="mt-1 text-sm leading-6 text-zinc-400">
                                {description}
                            </p>
                        )}
                    </div>

                    <span
                        className={cn(
                            "inline-flex shrink-0 items-center gap-1 rounded-full border px-2 py-1 text-[11px] font-medium",
                            statusClasses.badge
                        )}
                    >
                        {getStatusIcon(status)}
                        {getStatusLabel(status)}
                    </span>
                </div>

                {children && (
                    <div className="mt-3 rounded-[18px] border border-white/8 bg-black/20 p-3">
                        {children}
                    </div>
                )}
            </div>
        </div>
    );
}

ChainOfThoughtStep.displayName = "ChainOfThoughtStep";

export type ChainOfThoughtSearchResultsProps = ComponentPropsWithoutRef<"div">;

export function ChainOfThoughtSearchResults({
    className,
    ...props
}: ChainOfThoughtSearchResultsProps) {
    return (
        <div
            className={cn("flex flex-wrap gap-2", className)}
            {...props}
        />
    );
}

export type ChainOfThoughtSearchResultProps = ComponentPropsWithoutRef<"span">;

export function ChainOfThoughtSearchResult({
    className,
    children,
    ...props
}: ChainOfThoughtSearchResultProps) {
    return (
        <span
            className={cn(
                "inline-flex items-center gap-1.5 rounded-full border border-white/10 bg-white/[0.04] px-3 py-1.5 text-xs font-medium text-zinc-300",
                className
            )}
            {...props}
        >
            <Link2 className="size-3.5 text-zinc-500" />
            {children}
        </span>
    );
}

export interface ChainOfThoughtImageProps
    extends ComponentPropsWithoutRef<"div"> {
    caption?: string;
}

export function ChainOfThoughtImage({
    caption,
    className,
    children,
    ...props
}: ChainOfThoughtImageProps) {
    return (
        <div
            className={cn(
                "rounded-[20px] border border-white/10 bg-[linear-gradient(180deg,rgba(255,255,255,0.05),rgba(255,255,255,0.02))] p-3",
                className
            )}
            {...props}
        >
            <div className="overflow-hidden rounded-2xl border border-white/8 bg-black/20">
                {children}
            </div>
            {caption && (
                <div className="mt-3 flex items-start gap-2 text-xs leading-5 text-zinc-400">
                    <ImageIcon className="mt-0.5 size-3.5 shrink-0 text-zinc-500" />
                    <p>{caption}</p>
                </div>
            )}
        </div>
    );
}

ChainOfThought.displayName = "ChainOfThought";
ChainOfThoughtHeader.displayName = "ChainOfThoughtHeader";
ChainOfThoughtContent.displayName = "ChainOfThoughtContent";
ChainOfThoughtSearchResults.displayName = "ChainOfThoughtSearchResults";
ChainOfThoughtSearchResult.displayName = "ChainOfThoughtSearchResult";
ChainOfThoughtImage.displayName = "ChainOfThoughtImage";

export default ChainOfThought;