'use client';

import { useEffect, useRef } from 'react';
import { useCanvasStore } from '@/lib/store/canvasStore';
import { GenerationService } from '@/lib/services/generationService';
import { DefaultService } from '@/lib/api';
import { AppNodeType, NodeType, NodeStatus } from '../types/node';

const wait = (ms: number) => new Promise((r) => setTimeout(r, ms));

/**
 * 画布加载后恢复进行中的异步任务轮询（与各节点写入 taskId 的逻辑统一）。
 * 当节点带有 taskId 且 status 为 WORKING 时（例如刷新后从存储恢复），
 * 按节点类型自动轮询对应 task_id，完成后更新节点结果。
 * 支持：图片、视频、音乐、音频(TTS)、歌词。
 */
export function useResumeTaskPolling() {
    const nodes = useCanvasStore((s) => s.nodes);
    const setNodes = useCanvasStore((s) => s.setNodes);
    const pollingTaskIdsRef = useRef<Set<string>>(new Set());

    useEffect(() => {
        const working = nodes.filter(
            (n): n is AppNodeType =>
                n.type === 'appNode' &&
                (n.data as { status?: string }).status === NodeStatus.WORKING &&
                !!((n.data as { taskId?: string }).taskId)
        );

        working.forEach((node) => {
            const taskId = (node.data as { taskId?: string }).taskId!;
            if (pollingTaskIdsRef.current.has(taskId)) return;
            pollingTaskIdsRef.current.add(taskId);

            const nodeId = node.id;
            const type = (node.data as { type?: string }).type;
            const generationType = (node.data as { generationType?: string }).generationType;
            const model = (node.data as { model?: string }).model;

            if (type === NodeType.LYRICS_GENERATOR) {
                pollLyricsAndUpdate(taskId, nodeId, setNodes, pollingTaskIdsRef).catch(() => {});
                return;
            }

            if (type === NodeType.AUDIO_GENERATOR) {
                const taskType = generationType === 'music' ? 'music' : 'audio';
                GenerationService.pollTask(taskId, taskType, model, 120).then((result) => {
                    pollingTaskIdsRef.current.delete(taskId);
                    setNodes((nds) =>
                        nds.map((n) => {
                            if (n.id !== nodeId || n.type !== 'appNode') return n;
                            if (result.status === 'success' && result.urls?.[0]) {
                                return {
                                    ...n,
                                    data: {
                                        ...n.data,
                                        status: NodeStatus.SUCCESS,
                                        audioUri: result.urls[0],
                                        taskId: undefined,
                                        error: undefined,
                                    },
                                } as AppNodeType;
                            }
                            return {
                                ...n,
                                data: {
                                    ...n.data,
                                    status: NodeStatus.ERROR,
                                    error: result.error || '生成失败',
                                    taskId: undefined,
                                },
                            } as AppNodeType;
                        })
                    );
                });
                return;
            }

            if (type === NodeType.IMAGE_GENERATOR) {
                GenerationService.pollTask(taskId, 'image', model, 120).then((result) => {
                    pollingTaskIdsRef.current.delete(taskId);
                    setNodes((nds) =>
                        nds.map((n) => {
                            if (n.id !== nodeId || n.type !== 'appNode') return n;
                            if (result.status === 'success' && result.urls?.length) {
                                return {
                                    ...n,
                                    data: {
                                        ...n.data,
                                        status: NodeStatus.SUCCESS,
                                        image: result.urls[0],
                                        images: result.urls,
                                        taskId: undefined,
                                        error: undefined,
                                    },
                                } as AppNodeType;
                            }
                            return {
                                ...n,
                                data: {
                                    ...n.data,
                                    status: NodeStatus.ERROR,
                                    error: result.error || '图片生成失败',
                                    taskId: undefined,
                                },
                            } as AppNodeType;
                        })
                    );
                });
                return;
            }

            if (type === NodeType.VIDEO_GENERATOR) {
                GenerationService.pollTask(taskId, 'video', model, 120).then((result) => {
                    pollingTaskIdsRef.current.delete(taskId);
                    setNodes((nds) =>
                        nds.map((n) => {
                            if (n.id !== nodeId || n.type !== 'appNode') return n;
                            if (result.status === 'success' && result.urls?.length) {
                                return {
                                    ...n,
                                    data: {
                                        ...n.data,
                                        status: NodeStatus.SUCCESS,
                                        videoUri: result.urls[0],
                                        videoUris: result.urls,
                                        taskId: undefined,
                                        error: undefined,
                                    },
                                } as AppNodeType;
                            }
                            return {
                                ...n,
                                data: {
                                    ...n.data,
                                    status: NodeStatus.ERROR,
                                    error: result.error || '视频生成失败',
                                    taskId: undefined,
                                },
                            } as AppNodeType;
                        })
                    );
                });
            }
        });
    }, [nodes, setNodes]);
}

async function pollLyricsAndUpdate(
    taskId: string,
    nodeId: string,
    setNodes: (updater: (prev: AppNodeType[]) => AppNodeType[]) => void,
    pollingTaskIdsRef: React.MutableRefObject<Set<string>>
) {
    const maxAttempts = 120;
    const interval = 2000;
    let attempts = 0;

    while (attempts < maxAttempts) {
        try {
            const taskResponse = await DefaultService.getTaskStatus(taskId);
            const taskData = (taskResponse as any).data || {};
            const rawStatus = String(taskData.status || '').toLowerCase();
            const status = rawStatus === 'succeeded' ? 'success' : rawStatus;

            if (status === 'success') {
                const result = taskData.result || {};
                const lyrics = (result.lyrics || '').toString();
                if (lyrics.trim()) {
                    pollingTaskIdsRef.current.delete(taskId);
                    setNodes((nds) =>
                        nds.map((n) => {
                            if (n.id !== nodeId || n.type !== 'appNode') return n;
                            return {
                                ...n,
                                data: {
                                    ...n.data,
                                    status: NodeStatus.SUCCESS,
                                    lyrics,
                                    title: result.song_title ?? n.data.title,
                                    taskId: undefined,
                                    error: undefined,
                                },
                            } as AppNodeType;
                        })
                    );
                    return;
                }
            }

            if (status === 'failed' || status === 'cancelled' || status === 'canceled') {
                pollingTaskIdsRef.current.delete(taskId);
                setNodes((nds) =>
                    nds.map((n) => {
                        if (n.id !== nodeId || n.type !== 'appNode') return n;
                        return {
                            ...n,
                            data: {
                                ...n.data,
                                status: NodeStatus.ERROR,
                                error: taskData.error || '歌词生成失败',
                                taskId: undefined,
                            },
                        } as AppNodeType;
                    })
                );
                return;
            }
        } catch (_) {
            // ignore single request error, keep polling
        }
        await wait(interval);
        attempts += 1;
    }

    pollingTaskIdsRef.current.delete(taskId);
    setNodes((nds) =>
        nds.map((n) => {
            if (n.id !== nodeId || n.type !== 'appNode') return n;
            return {
                ...n,
                data: {
                    ...n.data,
                    status: NodeStatus.ERROR,
                    error: '歌词生成超时',
                    taskId: undefined,
                },
            } as AppNodeType;
        })
    );
}
