import { useCallback } from 'react';
import { 
    Connection, 
    Edge, 
    addEdge, 
    MarkerType,
} from '@xyflow/react';
import { AppNodeType, AppNodeData } from '../types/node';

// Define valid connection types
// Key: Source Node Type, Value: Array of valid Target Node Types
export const VALID_CONNECTIONS: Record<string, string[]> = {
    'PROMPT_INPUT': ['LYRICS_GENERATOR', 'IMAGE_GENERATOR', 'VIDEO_GENERATOR', 'AUDIO_GENERATOR', 'PROMPT_GENERATOR'],
    'LYRICS_GENERATOR': ['AUDIO_GENERATOR'],
    'PROMPT_GENERATOR': ['PROMPT_INPUT', 'LYRICS_GENERATOR', 'IMAGE_GENERATOR', 'VIDEO_GENERATOR'],
    'IMAGE_GENERATOR': ['IMAGE_GENERATOR', 'VIDEO_GENERATOR', 'IMAGE_EDITOR', 'VIDEO_ANALYZER'],
    'VIDEO_GENERATOR': ['VIDEO_GENERATOR', 'VIDEO_ANALYZER', 'IMAGE_EDITOR'],
    'AUDIO_GENERATOR': ['VIDEO_GENERATOR'],
    'IMAGE_EDITOR': ['IMAGE_GENERATOR', 'VIDEO_GENERATOR', 'VIDEO_ANALYZER'],
    'VIDEO_ANALYZER': ['PROMPT_GENERATOR', 'INFO_DISPLAY', 'IMAGE_GENERATOR', 'VIDEO_GENERATOR'], 
    'INFO_DISPLAY': ['IMAGE_GENERATOR', 'VIDEO_GENERATOR', 'IMAGE_EDITOR', 'VIDEO_ANALYZER', 'PROMPT_INPUT',"PROMPT_GENERATOR"] // Can act as source
};

export function useSocketInteraction(
    nodes: AppNodeType[],
    setEdges: React.Dispatch<React.SetStateAction<Edge[]>>
) {
    const isValidConnection = useCallback((connection: Connection | Edge) => {
        const sourceNode = nodes.find(n => n.id === connection.source);
        const targetNode = nodes.find(n => n.id === connection.target);

        if (!sourceNode || !targetNode) return false;

        // Ensure we are connecting AppNodes (not Groups) and they have the required data
        if (sourceNode.type !== 'appNode' || targetNode.type !== 'appNode') return false;

        const sourceType = sourceNode.data.type;
        const targetType = targetNode.data.type;

        // Prevent self-connection
        if (connection.source === connection.target) return false;

        // Check against valid connections map
        const validTargets = VALID_CONNECTIONS[sourceType];
        if (!validTargets || !validTargets.includes(targetType)) {
            // Optional: You could add a toast notification here explaining why it's invalid
            console.warn(`Invalid connection: ${sourceType} -> ${targetType}`);
            return false;
        }

        return true;
    }, [nodes]); // Depend on nodes to ensure we have latest data

    const onConnect = useCallback(
        (params: Connection) => {
            if (!isValidConnection(params)) return;

            setEdges((eds) => addEdge({ 
                ...params, 
                type: 'gradient',
                // animated: true, // Removed to match canvas solid gradient style
            }, eds));
        },
        [setEdges, isValidConnection]
    );

    return {
        onConnect,
        isValidConnection
    };
}
