import { useCallback, useState } from 'react';

type ExpandState = Set<number | string> | 'all';

export function useCardTableExpand(
    defaultExpanded: boolean | (number | string)[] = false,
) {
    const [expandedKeys, setExpandedKeys] = useState<ExpandState>(() => {
        if (defaultExpanded === true) return 'all';
        if (Array.isArray(defaultExpanded)) return new Set(defaultExpanded);
        return new Set();
    });

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

    const expandAll = useCallback((keys: (number | string)[]) => {
        setExpandedKeys(new Set(keys));
    }, []);

    const collapseAll = useCallback(() => {
        setExpandedKeys(new Set());
    }, []);

    const isExpanded = useCallback(
        (key: number | string): boolean => {
            if (expandedKeys === 'all') return true;
            return expandedKeys.has(key);
        },
        [expandedKeys],
    );

    return { expandedKeys, toggleExpand, expandAll, collapseAll, isExpanded };
}

export type { ExpandState };