import { useMemo } from 'react';
import { useNodeConnections, useNodesData } from '@xyflow/react';
import { useModelStore } from '@lib/store/modelStore';
import { AppNodeData, InputAsset } from '../types/node';
import { sortInputAssetsByTypeOrder } from '../utils';

function getVoiceGenderFromNode(audioModels: Record<string, { voices?: Array<{ id?: string; value?: string; gender?: string }> }>, model?: string, voiceId?: string): string | undefined {
    if (!model || !voiceId) return undefined;
    const voices = audioModels[model]?.voices;
    if (!voices?.length) return undefined;
    const v = voices.find((v: any) => (v.value || v.id) === voiceId);
    return v?.gender;
}

export function useNodeInput(nodeId: string) {
    const { audioModels } = useModelStore();
    const connections = useNodeConnections({ handleType: 'target' });
    const sourceIds = connections.map((c) => c.source);
    const nodesData = useNodesData(sourceIds);

    const inputs = useMemo(() => {
        let inputPrompt: string | undefined = undefined;
        let inputAssets: InputAsset[] = [];
        
        // useNodesData can return a single object or an array
        const sourceNodes = Array.isArray(nodesData) ? nodesData : (nodesData ? [nodesData] : []);

        sourceNodes.forEach((node) => {
            if (!node || !node.data) return;
            const data = node.data as AppNodeData;
            const label = data.title;
            
            // --- Extract Prompt ---
            let text = '';
            // For PROMPT_INPUT and PROMPT_GENERATOR, we use the 'prompt' field
            if (data.type === 'PROMPT_INPUT' || data.type === 'PROMPT_GENERATOR') {
                text = data.prompt || '';
            }
            if (data.type === 'LYRICS_GENERATOR') {
                text = data.lyrics || data.prompt || '';
            } 
            // For VIDEO_ANALYZER, we might want the analysis result, or the prompt if analysis is missing
            else if (data.type === 'VIDEO_ANALYZER') {
                text = data.analysis || data.prompt || '';
            }
            // For other nodes, maybe they pass through prompt? Or we don't extract text from them.
            
            if (text) {
                // If multiple inputs, join them with newlines
                // If inputPrompt was undefined, start with text
                inputPrompt = inputPrompt !== undefined ? `${inputPrompt}\n\n${text}` : text;
            }

            // --- Extract Assets ---
            const seenUrls = new Set<string>();

            // Single Image
            if (data.image && !seenUrls.has(data.image)) {
                inputAssets.push({ id: `${node.id}-img-main`, type: 'image', src: data.image, label });
                seenUrls.add(data.image);
            }
            // Multiple Images
            if (data.images && Array.isArray(data.images)) {
                data.images.forEach((img, idx) => {
                    if (img && !seenUrls.has(img)) {
                        inputAssets.push({ id: `${node.id}-img-${idx}`, type: 'image', src: img, label });
                        seenUrls.add(img);
                    }
                });
            }
            // Single Video
            if (data.videoUri && !seenUrls.has(data.videoUri)) {
                inputAssets.push({ id: `${node.id}-vid-main`, type: 'video', src: data.videoUri, label });
                seenUrls.add(data.videoUri);
            }
            // Multiple Videos
            if (data.videoUris && Array.isArray(data.videoUris)) {
                data.videoUris.forEach((vid, idx) => {
                    if (vid && !seenUrls.has(vid)) {
                        inputAssets.push({ id: `${node.id}-vid-${idx}`, type: 'video', src: vid, label });
                        seenUrls.add(vid);
                    }
                });
            }
            // Audio (e.g. AUDIO_GENERATOR 节点)：带上音色性别用于显示男女头像
            if (data.audioUri && !seenUrls.has(data.audioUri)) {
                const gender = getVoiceGenderFromNode(audioModels, data.model, data.voice);
                inputAssets.push({ id: `${node.id}-audio-main`, type: 'audio', src: data.audioUri, gender, label });
                seenUrls.add(data.audioUri);
            }
        });

        return { inputPrompt, inputAssets: sortInputAssetsByTypeOrder(inputAssets) };
    }, [nodesData, audioModels]);

    return inputs;
}
