import { useCallback } from 'react';
import { useReactFlow, Edge, Node } from '@xyflow/react';
import { AppNodeType, AppNodeData } from '../types/node';
import { useModelStore } from '@lib/store/modelStore';
import { useCanvasActionContext } from '../CanvasContext';
import { useProjectStore } from '@lib/store/project';
import { getExecutionStrategy, ExecutionContext } from '../services/execution';

// Helper to get input nodes based on edges
const getInputNodes = (nodeId: string, nodes: AppNodeType[], edges: Edge[]): Node<AppNodeData, 'appNode'>[] => {
    const inputIds = edges
        .filter(e => e.target === nodeId)
        .map(e => e.source);
    // Only return AppNodes (filter out Groups or other types if any)
    return nodes.filter((n): n is Node<AppNodeData, 'appNode'> => inputIds.includes(n.id) && n.type === 'appNode');
};

export function useNodeExecution() {
    const { getNodes, setNodes, getEdges, setEdges, updateNodeData } = useReactFlow<AppNodeType>();
    const projectId = useProjectStore((s) => s.id);
    const { addHistoryItem } = useCanvasActionContext();
    const { getDefaultModelId } = useModelStore();

    const getLLMModelDefault = useCallback(() => getDefaultModelId('llm'), [getDefaultModelId]);
    const getImageModelDefault = useCallback(() => getDefaultModelId('image'), [getDefaultModelId]);
    const getVideoModelDefault = useCallback(() => getDefaultModelId('video'), [getDefaultModelId]);
    const getAudioModelDefault = useCallback(() => getDefaultModelId('audio'), [getDefaultModelId]);
    const getMusicModelDefault = useCallback(() => getDefaultModelId('music'), [getDefaultModelId]);

    const executeNode = useCallback(async (id: string, promptOverride?: string) => {
        const nodes = getNodes();
        const edges = getEdges();
        const node = nodes.find(n => n.id === id);
        
        if (!node || node.type !== 'appNode') return;

        // Update status to WORKING
        updateNodeData(id, { status: 'WORKING', error: undefined });

        try {
            const inputs = getInputNodes(id, nodes, edges);
            
            // Gather upstream texts
            const upstreamTexts = inputs.map(n => {
                if (n.type !== 'appNode') return null;
                const data = n.data;
                
                if (data.type === 'PROMPT_INPUT') return data.prompt;
                if (data.type === 'PROMPT_GENERATOR') return data.prompt;
                if (data.type === 'LYRICS_GENERATOR') return data.lyrics || data.prompt;
                if (data.type === 'VIDEO_ANALYZER') return data.analysis;
                // INFO_DISPLAY nodes: Only used for visual connection indication
                // They don't pass any information to downstream nodes
                // Downstream nodes will use their own prompts
                if (data.type === 'INFO_DISPLAY') {
                    return null;
                }
                return null;
            }).filter((t): t is string => !!t && t.trim().length > 0);

            // Construct prompt:
            // - Music node keeps its own prompt (style/scene description)
            // - Other nodes can still consume upstream text as their main prompt
            let prompt = promptOverride || '';
            const isMusicNode = node.data.type === 'AUDIO_GENERATOR' && node.data.generationType === 'music';
            if (isMusicNode) {
                prompt = prompt || node.data.prompt || '';
            } else if (upstreamTexts.length > 0) {
                prompt = upstreamTexts.join('\n');
            } else {
                prompt = node.data.prompt || '';
            }

            // Prepend style prompt if available (applied at execution time)
            if (node.data.stylePrompt) {
                prompt = prompt ? `${node.data.stylePrompt}, ${prompt}` : node.data.stylePrompt;
            }

            // Create Execution Context
            const context: ExecutionContext = {
                nodeId: id,
                nodes,
                edges,
                projectId: projectId || undefined,
                prompt,
                upstreamTexts,
                inputs: inputs as any[],
                node: node as any,
                updateNodeData,
                setNodes,
                setEdges,
                addHistoryItem,
                getLLMModelDefault,
                getImageModelDefault,
                getVideoModelDefault,
                getAudioModelDefault,
                getMusicModelDefault
            };

            // Get and execute strategy
            const strategy = getExecutionStrategy(node.data.type);
            await strategy.execute(context);

        } catch (error: any) {
            console.error('[useNodeExecution] Execution Failed:', error);
            updateNodeData(id, { 
                taskId: undefined,
                status: 'ERROR', 
                error: error.message || 'Execution failed' 
            });
        }
    }, [getNodes, getEdges, updateNodeData, getLLMModelDefault, getImageModelDefault, getVideoModelDefault, getAudioModelDefault, getMusicModelDefault, addHistoryItem, setEdges, setNodes, projectId]);

    return { executeNode };
}
