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

interface ArmBotStatus extends RoboticArmSchemaEvent {
    productId: string;
    isActive: boolean;
    idleTime: number;
    activeTime: number;
    currentAction?: {
        actionUuid: string;
        actionName: string;
    };
}

interface ArmBotStore {
    armBots: Record<string, ArmBotStatus>;

    // ArmBot actions
    addArmBot: (productId: string, event: RoboticArmSchemaEvent) => void;
    removeArmBot: (modelUuid: string) => void;
    updateArmBot: (
        modelUuid: string,
        updates: Partial<Omit<ArmBotStatus, 'modelUuid' | 'productId'>>
    ) => void;

    // Action management
    startAction: (modelUuid: string, actionUuid: string) => void;
    completeAction: (modelUuid: string) => void;
    cancelAction: (modelUuid: string) => void;

    // Status updates
    setArmBotActive: (modelUuid: string, isActive: boolean) => void;

    // Time tracking
    incrementActiveTime: (modelUuid: string, incrementBy: number) => void;
    incrementIdleTime: (modelUuid: string, incrementBy: number) => void;

    // Helper functions
    getArmBotById: (modelUuid: string) => ArmBotStatus | undefined;
    getArmBotsByProduct: (productId: string) => ArmBotStatus[];
    getActiveArmBots: () => ArmBotStatus[];
    getIdleArmBots: () => ArmBotStatus[];
    getArmBotsByCurrentAction: (actionUuid: string) => ArmBotStatus[];
}

export const useArmBotStore = create<ArmBotStore>()(
    immer((set, get) => ({
        armBots: {},

        // ArmBot actions
        addArmBot: (productId, event) => {
            set((state) => {
                state.armBots[event.modelUuid] = {
                    ...event,
                    productId,
                    isActive: false,
                    idleTime: 0,
                    activeTime: 0,
                    state: 'idle'
                };
            });
        },

        removeArmBot: (modelUuid) => {
            set((state) => {
                delete state.armBots[modelUuid];
            });
        },

        updateArmBot: (modelUuid, updates) => {
            set((state) => {
                const armBot = state.armBots[modelUuid];
                if (armBot) {
                    Object.assign(armBot, updates);
                }
            });
        },

        // Action management
        startAction: (modelUuid, actionUuid) => {
            set((state) => {
                const armBot = state.armBots[modelUuid];
                if (armBot) {
                    const action = armBot.point.actions.find(a => a.actionUuid === actionUuid);
                    if (action) {
                        armBot.currentAction = {
                            actionUuid: action.actionUuid,
                            actionName: action.actionName
                        };
                        armBot.isActive = true;
                    }
                }
            });
        },

        completeAction: (modelUuid) => {
            set((state) => {
                const armBot = state.armBots[modelUuid];
                if (armBot && armBot.currentAction) {
                    armBot.currentAction = undefined;
                    armBot.isActive = false;
                }
            });
        },

        cancelAction: (modelUuid) => {
            set((state) => {
                const armBot = state.armBots[modelUuid];
                if (armBot) {
                    armBot.currentAction = undefined;
                    armBot.isActive = false;
                }
            });
        },

        // Status updates
        setArmBotActive: (modelUuid, isActive) => {
            set((state) => {
                const armBot = state.armBots[modelUuid];
                if (armBot) {
                    armBot.isActive = isActive;
                }
            });
        },

        // Time tracking
        incrementActiveTime: (modelUuid, incrementBy) => {
            set((state) => {
                const armBot = state.armBots[modelUuid];
                if (armBot) {
                    armBot.activeTime += incrementBy;
                }
            });
        },

        incrementIdleTime: (modelUuid, incrementBy) => {
            set((state) => {
                const armBot = state.armBots[modelUuid];
                if (armBot) {
                    armBot.idleTime += incrementBy;
                }
            });
        },

        // Helper functions
        getArmBotById: (modelUuid) => {
            return get().armBots[modelUuid];
        },

        getArmBotsByProduct: (productId) => {
            return Object.values(get().armBots).filter(
                a => a.productId === productId
            );
        },

        getActiveArmBots: () => {
            return Object.values(get().armBots).filter(a => a.isActive);
        },

        getIdleArmBots: () => {
            return Object.values(get().armBots).filter(
                a => !a.isActive && a.state === 'idle'
            );
        },

        getArmBotsByCurrentAction: (actionUuid) => {
            return Object.values(get().armBots).filter(
                a => a.currentAction?.actionUuid === actionUuid
            );
        }
    }))
);