import { useState, useEffect, useCallback } from 'react';
import { TaskInfo } from 'shared/domain/librarian/task';
import { useQuery } from '@tanstack/react-query';
import { taskApi } from './api';

export function useTaskStatus(taskId: string | null, taskType: string) {
    const [taskInfo, setTaskInfo] = useState<TaskInfo | null>(null);
    const [error, setError] = useState<Error | null>(null);

    const handleUpdate = useCallback((info: TaskInfo) => {
        setTaskInfo(info);
        setError(null);
    }, []);

    const handleError = useCallback((err: any) => {
        console.error('TaskWebSocket Error:', err);
        const newError =
            err instanceof Error
                ? err
                : new Error(
                      String(err.message || 'WebSocket connection error')
                  );
        setError(newError);
        setTaskInfo(null);
    }, []);

    const { data, isError, error: queryError } = useQuery(
        ['taskStatus', taskId, taskType],
        () => taskApi.getTaskStatusApi(taskType, taskId!),
        {
            enabled: !!taskId,
            onSuccess: handleUpdate,
            onError: handleError,
        }
    );

    useEffect(() => {
        if (isError) {
            handleError(queryError);
        }
    }, [isError, queryError, handleError]);

    return { taskInfo: data || taskInfo, error };
}