Dwinzo_Demo/app/src/store/builder/useAssetGroupStore.ts

import { create } from "zustand";
import { immer } from "zustand/middleware/immer";

interface AssetGroupStore {
    assetGroups: AssetGroup[];
    selectedGroups: string[]; // Array of groupUuids
    groupHierarchy: AssetGroupHierarchy;

    // Group CRUD operations
    addGroup: (group: AssetGroup) => { createdGroups: AssetGroup[]; updatedGroups: AssetGroup[] };
    removeGroup: (groupUuid: string) => void;
    clearGroups: () => void;
    setGroups: (groups: AssetGroup[]) => void;

    // Group selection
    setSelectedGroups: (groupUuids: string[]) => void;
    addSelectedGroup: (groupUuid: string) => void;
    removeSelectedGroup: (groupUuid: string) => void;
    toggleSelectedGroup: (groupUuid: string) => void;
    clearSelectedGroups: () => void;
    hasSelectedGroup: (groupUuid: string) => boolean;

    // Group children management
    addChildToGroup: (groupUuid: string, child: { type: "Asset" | "Group"; childrenUuid: string }) => { updatedGroups: AssetGroup[] };
    removeChildFromGroup: (groupUuid: string, childUuid: string) => void;
    getGroupChildren: (groupUuid: string) => { type: "Asset" | "Group"; childrenUuid: string }[];

    // Group properties
    setGroupName: (groupUuid: string, newName: string) => void;
    setGroupVisibility: (groupUuid: string, isVisible: boolean) => void;
    setGroupLock: (groupUuid: string, isLocked: boolean) => void;
    toggleGroupVisibility: (groupUuid: string) => void;
    setGroupExpanded: (groupUuid: string, isExpanded: boolean) => void;

    // Hierarchy operations
    buildHierarchy: (assets: Assets, groups: AssetGroup[]) => AssetGroupHierarchy;
    flattenHierarchy: (hierarchy: AssetGroupHierarchy) => { assets: Asset[]; groups: AssetGroup[] };
    getGroupHierarchy: (groupUuid: string) => AssetGroupHierarchyNode | null;
    getFlatGroupAssets: (groupUuid: string, assets: Assets) => Asset[];
    getFlatGroupChildren: (groupUuid: string) => string[]; // Returns all child uuids (both assets and groups)

    // Helper functions
    getGroupById: (groupUuid: string) => AssetGroup | undefined;
    getGroupsContainingAsset: (assetUuid: string) => AssetGroup[];
    getGroupsContainingGroup: (childGroupUuid: string) => AssetGroup[];
    getParentGroup: (item: AssetGroupChild) => string | null;
    hasGroup: (groupUuid: string) => boolean;
    isGroup: (item: AssetGroupChild) => item is AssetGroupHierarchyNode;
    isEmptyGroup: (groupUuid: string) => boolean;
}

export const createAssetGroupStore = () => {
    return create<AssetGroupStore>()(
        immer((set, get) => ({
            assetGroups: [],
            selectedGroups: [],
            groupHierarchy: [],

            // Group CRUD operations
            addGroup: (group) => {
                const result: { createdGroups: AssetGroup[]; updatedGroups: AssetGroup[] } = {
                    createdGroups: [],
                    updatedGroups: [],
                };

                set((state) => {
                    // Check if group already exists
                    if (state.assetGroups.some((g) => g.groupUuid === group.groupUuid)) {
                        return;
                    }

                    // Find all asset children in the new group
                    const assetChildren = group.children.filter((child) => child.type === "Asset");
                    const assetUuids = new Set(assetChildren.map((child) => child.childrenUuid));

                    // Remove these assets from existing groups and track updated groups
                    const updatedGroups: AssetGroup[] = [];

                    state.assetGroups.forEach((existingGroup) => {
                        const originalLength = existingGroup.children.length;
                        existingGroup.children = existingGroup.children.filter((child) => !(child.type === "Asset" && assetUuids.has(child.childrenUuid)));

                        // If group was modified, add to updated groups
                        if (existingGroup.children.length !== originalLength) {
                            updatedGroups.push({ ...existingGroup });
                        }
                    });

                    // Add the new group
                    state.assetGroups.push(group);
                    result.createdGroups.push({ ...group });
                    result.updatedGroups = updatedGroups;
                });

                return result;
            },

            removeGroup: (groupUuid) => {
                set((state) => {
                    // First remove this group from any parent groups
                    state.assetGroups.forEach((group) => {
                        group.children = group.children.filter((child) => !(child.type === "Group" && child.childrenUuid === groupUuid));
                    });

                    // Then remove the group itself
                    state.assetGroups = state.assetGroups.filter((g) => g.groupUuid !== groupUuid);

                    // Remove from selected groups
                    state.selectedGroups = state.selectedGroups.filter((uuid) => uuid !== groupUuid);
                });
            },

            clearGroups: () => {
                set((state) => {
                    state.assetGroups = [];
                    state.selectedGroups = [];
                    state.groupHierarchy = [];
                });
            },

            setGroups: (groups) => {
                set((state) => {
                    state.assetGroups = groups;
                });
            },

            // Group selection
            setSelectedGroups: (groupUuids) => {
                set((state) => {
                    state.selectedGroups = groupUuids;
                });
            },

            addSelectedGroup: (groupUuid) => {
                set((state) => {
                    if (!state.selectedGroups.includes(groupUuid)) {
                        state.selectedGroups.push(groupUuid);
                    }
                });
            },

            removeSelectedGroup: (groupUuid) => {
                set((state) => {
                    state.selectedGroups = state.selectedGroups.filter((uuid) => uuid !== groupUuid);
                });
            },

            toggleSelectedGroup: (groupUuid) => {
                set((state) => {
                    const exists = state.selectedGroups.includes(groupUuid);
                    if (exists) {
                        state.selectedGroups = state.selectedGroups.filter((uuid) => uuid !== groupUuid);
                    } else {
                        state.selectedGroups.push(groupUuid);
                    }
                });
            },

            clearSelectedGroups: () => {
                set((state) => {
                    state.selectedGroups = [];
                });
            },

            hasSelectedGroup: (groupUuid: string) => {
                return get().selectedGroups.includes(groupUuid);
            },

            // Group children management
            addChildToGroup: (groupUuid, child) => {
                const result: { updatedGroups: AssetGroup[] } = { updatedGroups: [] };

                set((state) => {
                    const targetGroup = state.assetGroups.find((g) => g.groupUuid === groupUuid);
                    if (!targetGroup) return;

                    const updatedGroups: AssetGroup[] = [];

                    // 1️⃣ Remove the child from any other groups (to maintain single-parent rule)
                    state.assetGroups.forEach((group) => {
                        if (group.groupUuid === groupUuid) return; // skip target group

                        const originalLength = group.children.length;
                        group.children = group.children.filter((c) => c.childrenUuid !== child.childrenUuid);

                        if (group.children.length !== originalLength) {
                            updatedGroups.push({ ...group });
                        }
                    });

                    // 2️⃣ Add the child to the target group (if not already present)
                    if (!targetGroup.children.some((c) => c.childrenUuid === child.childrenUuid)) {
                        targetGroup.children.push(child);
                        updatedGroups.push({ ...targetGroup });
                    }

                    // 3️⃣ Rebuild hierarchy after modification
                    state.groupHierarchy = get().buildHierarchy([], state.assetGroups);

                    result.updatedGroups = updatedGroups;
                });

                return result;
            },

            removeChildFromGroup: (groupUuid, childUuid) => {
                set((state) => {
                    const group = state.assetGroups.find((g) => g.groupUuid === groupUuid);
                    if (group) {
                        group.children = group.children.filter((child) => child.childrenUuid !== childUuid);
                        state.groupHierarchy = get().buildHierarchy([], state.assetGroups);
                    }
                });
            },

            getGroupChildren: (groupUuid) => {
                const group = get().assetGroups.find((g) => g.groupUuid === groupUuid);
                return group?.children || [];
            },

            // Group properties
            setGroupName: (groupUuid, newName) => {
                set((state) => {
                    const group = state.assetGroups.find((g) => g.groupUuid === groupUuid);
                    if (group) {
                        group.groupName = newName;
                    }
                });
            },

            setGroupVisibility: (groupUuid, isVisible) => {
                set((state) => {
                    const group = state.assetGroups.find((g) => g.groupUuid === groupUuid);
                    if (group) {
                        group.isVisible = isVisible;
                    }
                });
            },

            setGroupLock: (groupUuid, isLocked) => {
                set((state) => {
                    const group = state.assetGroups.find((g) => g.groupUuid === groupUuid);
                    if (group) {
                        group.isLocked = isLocked;
                    }
                });
            },

            toggleGroupVisibility: (groupUuid) => {
                set((state) => {
                    const group = state.assetGroups.find((g) => g.groupUuid === groupUuid);
                    if (group) {
                        group.isVisible = !group.isVisible;
                    }
                });
            },

            setGroupExpanded: (groupUuid, isExpanded) => {
                set((state) => {
                    const group = state.assetGroups.find((g) => g.groupUuid === groupUuid);
                    if (group) {
                        group.isExpanded = isExpanded;
                    }
                });
            },

            // Hierarchy operations
            buildHierarchy: (assets: Assets, groups: AssetGroup[]): AssetGroupHierarchy => {
                const assetMap = new Map(assets.map((asset) => [asset.modelUuid, asset]));
                const groupMap = new Map(groups.map((group) => [group.groupUuid, group]));

                const buildNode = (group: AssetGroup): AssetGroupHierarchyNode => {
                    const children: AssetGroupChild[] = [];

                    group.children.forEach((child) => {
                        if (child.type === "Asset") {
                            const asset = assetMap.get(child.childrenUuid);
                            if (asset) {
                                children.push(asset);
                                // Remove from assetMap so we know it's been processed
                                assetMap.delete(child.childrenUuid);
                            }
                        } else if (child.type === "Group") {
                            const childGroup = groupMap.get(child.childrenUuid);
                            if (childGroup) {
                                children.push(buildNode(childGroup));
                            }
                        }
                    });

                    return {
                        groupUuid: group.groupUuid,
                        groupName: group.groupName,
                        isVisible: group.isVisible,
                        isLocked: group.isLocked,
                        isExpanded: group.isExpanded,
                        children,
                    };
                };

                // Find root groups (groups that are not children of any other group)
                const childGroupUuids = new Set();
                groups.forEach((group) => {
                    group.children.forEach((child) => {
                        if (child.type === "Group") {
                            childGroupUuids.add(child.childrenUuid);
                        }
                    });
                });

                const rootGroups = groups.filter((group) => !childGroupUuids.has(group.groupUuid));

                // Build hierarchy starting from root groups
                const hierarchy: AssetGroupHierarchy = rootGroups.map(buildNode);

                // Add remaining assets that are not in any group
                const ungroupedAssets: Asset[] = [];
                assetMap.forEach((asset) => {
                    ungroupedAssets.push(asset);
                });

                const finalHierarchy = [...hierarchy, ...ungroupedAssets];

                set((state) => {
                    state.groupHierarchy = finalHierarchy;
                });

                return finalHierarchy;
            },

            flattenHierarchy: (hierarchy: AssetGroupHierarchy) => {
                const assets: Asset[] = [];
                const groups: AssetGroup[] = [];
                const processedGroups = new Set<string>();

                const processNode = (node: AssetGroupChild) => {
                    if ("modelUuid" in node) {
                        // It's an Asset
                        assets.push(node);
                    } else {
                        // It's an AssetGroupHierarchyNode
                        if (!processedGroups.has(node.groupUuid)) {
                            groups.push({
                                groupUuid: node.groupUuid,
                                groupName: node.groupName,
                                isVisible: node.isVisible,
                                isLocked: node.isLocked,
                                isExpanded: node.isExpanded,
                                children: node.children.map((child) =>
                                    "modelUuid" in child ? { type: "Asset" as const, childrenUuid: child.modelUuid } : { type: "Group" as const, childrenUuid: child.groupUuid }
                                ),
                            });
                            processedGroups.add(node.groupUuid);
                        }

                        node.children.forEach(processNode);
                    }
                };

                hierarchy.forEach(processNode);
                return { assets, groups };
            },

            getGroupHierarchy: (groupUuid) => {
                const hierarchy = get().groupHierarchy;

                const findGroup = (nodes: AssetGroupHierarchy): AssetGroupHierarchyNode | null => {
                    for (const node of nodes) {
                        if ("groupUuid" in node && node.groupUuid === groupUuid) {
                            return node;
                        }
                        if ("children" in node) {
                            const found = findGroup(node.children);
                            if (found) return found;
                        }
                    }
                    return null;
                };

                return findGroup(hierarchy);
            },

            getFlatGroupAssets: (groupUuid, assets) => {
                const groupHierarchy = get().getGroupHierarchy(groupUuid);
                if (!groupHierarchy) return [];

                const flatAssets: Asset[] = [];
                const assetMap = new Map(assets.map((asset) => [asset.modelUuid, asset]));

                const collectAssets = (nodes: AssetGroupChild[]) => {
                    nodes.forEach((node) => {
                        if ("modelUuid" in node) {
                            const asset = assetMap.get(node.modelUuid);
                            if (asset) flatAssets.push(asset);
                        } else {
                            collectAssets(node.children);
                        }
                    });
                };

                collectAssets(groupHierarchy.children);
                return flatAssets;
            },

            getFlatGroupChildren: (groupUuid) => {
                const group = get().assetGroups.find((g) => g.groupUuid === groupUuid);
                if (!group) return [];

                const allChildren: string[] = [];

                const collectChildren = (children: { type: "Asset" | "Group"; childrenUuid: string }[]) => {
                    children.forEach((child) => {
                        allChildren.push(child.childrenUuid);
                        if (child.type === "Group") {
                            const childGroup = get().assetGroups.find((g) => g.groupUuid === child.childrenUuid);
                            if (childGroup) {
                                collectChildren(childGroup.children);
                            }
                        }
                    });
                };

                collectChildren(group.children);
                return allChildren;
            },

            // Helper functions
            getGroupById: (groupUuid) => {
                return get().assetGroups.find((g) => g.groupUuid === groupUuid);
            },

            getGroupsContainingAsset: (assetUuid) => {
                return get().assetGroups.filter((group) => group.children.some((child) => child.type === "Asset" && child.childrenUuid === assetUuid));
            },

            getGroupsContainingGroup: (childGroupUuid) => {
                return get().assetGroups.filter((group) => group.children.some((child) => child.type === "Group" && child.childrenUuid === childGroupUuid));
            },

            getParentGroup: (item) => {
                if (get().isGroup(item)) {
                    const parents = get().getGroupsContainingGroup(item.groupUuid);
                    return parents.length > 0 ? parents[0].groupUuid : null;
                } else {
                    const parents = get().getGroupsContainingAsset(item.modelUuid);
                    return parents.length > 0 ? parents[0].groupUuid : null;
                }
            },

            hasGroup: (groupUuid) => {
                return get().assetGroups.some((g) => g.groupUuid === groupUuid);
            },

            isGroup: (item): item is AssetGroupHierarchyNode => {
                return "children" in item;
            },

            isEmptyGroup: (groupUuid) => {
                const group = get().assetGroups.find((g) => g.groupUuid === groupUuid);
                return !group || group.children.length === 0;
            },
        }))
    );
};

export type AssetGroupStoreType = ReturnType<typeof createAssetGroupStore>;