wrongmove/frontend/src/components/TaskIndicator.tsx

import { getUser } from '@/auth/authService';
import { getStoredPasskeyUser } from '@/auth/passkeyService';
import { fromOidcUser, type AuthUser } from '@/auth/types';
import { POLLING_INTERVALS } from '@/constants';
import { fetchTaskStatus, cancelTask, clearAllTasks } from '@/services';
import { TaskStatus, type TaskResult, type TaskState } from '@/types';
import { useEffect, useState, useRef, useMemo } from 'react';
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from './ui/tooltip';
import { Button } from './ui/button';
import { Loader2, CheckCircle2, XCircle, X, Trash2 } from 'lucide-react';
import { TaskProgressDrawer } from './TaskProgressDrawer';

interface TaskIndicatorProps {
    taskID: string | null;
    onTaskCancelled?: () => void;
    onTaskCompleted?: () => void;
    wsTasks?: Record<string, TaskState>;
    wsConnected?: boolean;
}

/** Convert a TaskState (from WS) into a TaskResult (for the drawer). */
function taskStateToResult(ts: TaskState): TaskResult {
    return {
        progress: ts.progress ?? 0,
        processed: ts.processed,
        total: ts.total,
        phase: ts.phase,
        message: ts.message,
        subqueries_probed: ts.subqueries_probed,
        subqueries_initial: ts.subqueries_initial,
        estimated_results: ts.estimated_results,
        subqueries_total: ts.subqueries_total,
        subqueries_completed: ts.subqueries_completed,
        ids_collected: ts.ids_collected,
        pages_fetched: ts.pages_fetched,
        fetching_done: ts.fetching_done,
        details_fetched: ts.details_fetched,
        images_downloaded: ts.images_downloaded,
        ocr_completed: ts.ocr_completed,
        failed: ts.failed,
        elapsed_seconds: ts.elapsed_seconds,
        rate_per_second: ts.rate_per_second,
        eta_seconds: ts.eta_seconds,
        logs: ts.logs,
    };
}

function isTerminalStatus(status: string): boolean {
    return status === 'SUCCESS' || status === 'FAILURE' || status === 'REVOKED';
}

export function TaskIndicator({
    taskID,
    onTaskCancelled,
    onTaskCompleted,
    wsTasks,
    wsConnected,
}: TaskIndicatorProps) {
    const [user, setUser] = useState<AuthUser | null>(null);
    const [progressPercentage, setProgressPercentage] = useState<number>(0);
    const [processed, setProcessed] = useState<number | null>(null);
    const [total, setTotal] = useState<number | null>(null);
    const [taskStatus, setTaskStatus] = useState<TaskStatus | null>(null);
    const [taskResult, setTaskResult] = useState<TaskResult | null>(null);
    const [isCancelling, setIsCancelling] = useState(false);
    const [isClearing, setIsClearing] = useState(false);
    const [drawerOpen, setDrawerOpen] = useState(false);
    const [selectedTaskId, setSelectedTaskId] = useState<string | null>(null);

    // Prevents WS effect from overwriting local cancel/clear state before
    // the parent's setTaskID(null) propagates. Reset when taskID changes.
    const cancelledRef = useRef(false);

    const onTaskCompletedRef = useRef(onTaskCompleted);
    useEffect(() => {
        onTaskCompletedRef.current = onTaskCompleted;
    }, [onTaskCompleted]);

    useEffect(() => {
        const passkeyUser = getStoredPasskeyUser();
        if (passkeyUser) {
            setUser(passkeyUser);
        } else {
            getUser().then((oidcUser) => {
                if (oidcUser) setUser(fromOidcUser(oidcUser));
            });
        }
    }, []);

    // Track the currently-viewed task in the drawer; default to the externally-provided taskID
    useEffect(() => {
        if (taskID) {
            setSelectedTaskId(taskID);
            cancelledRef.current = false; // new task, reset cancelled guard
        }
    }, [taskID]);

    // Count active (non-terminal) tasks from WS
    const activeWsTaskCount = useMemo(() => {
        if (!wsTasks) return 0;
        return Object.values(wsTasks).filter(
            (t) => !isTerminalStatus(t.status),
        ).length;
    }, [wsTasks]);

    // ----- WebSocket-driven state updates -----
    // When wsConnected, derive taskStatus/taskResult/progress from wsTasks
    useEffect(() => {
        if (!wsConnected || !wsTasks || !taskID) return;
        // Don't let WS overwrite local cancel/clear state
        if (cancelledRef.current) return;
        const wsTask = wsTasks[taskID];
        if (!wsTask) return;

        const status = wsTask.status as TaskStatus;
        setTaskStatus(status);

        if (wsTask.phase) {
            setTaskResult(taskStateToResult(wsTask));
        }
        if (wsTask.progress !== undefined) {
            setProgressPercentage(wsTask.progress * 100);
        }
        if (wsTask.processed !== undefined) {
            setProcessed(wsTask.processed);
        }
        if (wsTask.total !== undefined) {
            setTotal(wsTask.total);
        }

        if (status === TaskStatus.SUCCESS) {
            setProgressPercentage(100);
            onTaskCompletedRef.current?.();
        }
    }, [wsConnected, wsTasks, taskID]);

    // ----- Polling fallback (only when WS is not connected) -----
    useEffect(() => {
        // If WS is connected, skip polling
        if (wsConnected) return;

        if (!user || !taskID) {
            setTaskStatus(null);
            setTaskResult(null);
            return;
        }

        // Reset state for new task
        setTaskStatus(TaskStatus.PENDING);
        setProgressPercentage(0);
        setProcessed(null);
        setTotal(null);
        setTaskResult(null);

        const pollTaskStatus = async () => {
            try {
                const data = await fetchTaskStatus(user, taskID);
                const status = data.status as TaskStatus;
                setTaskStatus(status);

                if (status === TaskStatus.SUCCESS) {
                    setProgressPercentage(100);
                    if (data.result) {
                        try {
                            const parsedResult: TaskResult = JSON.parse(data.result);
                            if (parsedResult.phase) {
                                setTaskResult(parsedResult);
                            }
                        } catch {
                            // Ignore parsing errors
                        }
                    }
                    onTaskCompletedRef.current?.();
                    return true;
                }

                if (status === TaskStatus.FAILURE || status === TaskStatus.REVOKED) {
                    return true;
                }

                if (data.result) {
                    try {
                        const parsedResult: TaskResult = JSON.parse(data.result);
                        if (parsedResult.phase) {
                            setTaskResult(parsedResult);
                        }
                        if (parsedResult.progress !== undefined) {
                            setProgressPercentage(parsedResult.progress * 100);
                        }
                        if (parsedResult.processed !== undefined) {
                            setProcessed(parsedResult.processed);
                        }
                        if (parsedResult.total !== undefined) {
                            setTotal(parsedResult.total);
                        }
                    } catch {
                        // Ignore parsing errors
                    }
                }
                return false;
            } catch {
                setTaskStatus(TaskStatus.FAILURE);
                return true;
            }
        };

        pollTaskStatus();

        const interval = setInterval(async () => {
            const shouldStop = await pollTaskStatus();
            if (shouldStop) {
                clearInterval(interval);
            }
        }, POLLING_INTERVALS.TASK_STATUS_MS);

        return () => clearInterval(interval);
    }, [taskID, user, wsConnected]);

    const handleCancel = async () => {
        if (!user || !taskID || isCancelling) return;

        setIsCancelling(true);
        try {
            const result = await cancelTask(user, taskID);
            if (result.success) {
                cancelledRef.current = true;
                setTaskStatus(TaskStatus.REVOKED);
                onTaskCancelled?.();
            }
        } catch {
            // Ignore cancel errors
        } finally {
            setIsCancelling(false);
        }
    };

    const handleClearAll = async () => {
        if (!user || isClearing) return;

        setIsClearing(true);
        try {
            const result = await clearAllTasks(user);
            if (result.success) {
                cancelledRef.current = true;
                setTaskStatus(null);
                setTaskResult(null);
                onTaskCancelled?.();
            }
        } catch {
            // Ignore clear errors
        } finally {
            setIsClearing(false);
        }
    };

    if (!taskID || !taskStatus) {
        return null;
    }

    const isInProgress = taskStatus !== TaskStatus.SUCCESS &&
        taskStatus !== TaskStatus.FAILURE &&
        taskStatus !== TaskStatus.REVOKED;

    const getStatusIcon = () => {
        if (isInProgress) {
            return <Loader2 className="h-3.5 w-3.5 animate-spin text-blue-500" />;
        }
        if (taskStatus === TaskStatus.SUCCESS) {
            return <CheckCircle2 className="h-3.5 w-3.5 text-green-500" />;
        }
        return <XCircle className="h-3.5 w-3.5 text-red-500" />;
    };

    const getProgressText = () => {
        if (processed !== null && total !== null && total > 0) {
            return `${processed} / ${total}`;
        }
        if (taskResult?.phase && taskResult.phase !== 'processing') {
            const phaseLabels: Record<string, string> = {
                splitting: 'Splitting',
                splitting_complete: 'Split done',
                fetching: 'Fetching',
            };
            return phaseLabels[taskResult.phase] ?? `${Math.round(progressPercentage)}%`;
        }
        return `${Math.round(progressPercentage)}%`;
    };

    const getTooltipContent = () => {
        if (isInProgress) {
            if (processed !== null && total !== null && total > 0) {
                return `Processing: ${processed} / ${total} listings (${Math.round(progressPercentage)}%) — click for details`;
            }
            return `Task running: ${getProgressText()} — click for details`;
        }
        if (taskStatus === TaskStatus.SUCCESS) {
            return 'Task completed successfully — click for details';
        }
        if (taskStatus === TaskStatus.REVOKED) {
            return 'Task was cancelled';
        }
        return 'Task failed';
    };

    return (
        <TooltipProvider>
            <div className="flex items-center gap-2">
                <Tooltip>
                    <TooltipTrigger asChild>
                        <div
                            className="flex items-center gap-2 cursor-pointer"
                            onClick={() => setDrawerOpen(true)}
                        >
                            {getStatusIcon()}
                            {isInProgress && (
                                <div className="flex items-center gap-2">
                                    <div className="w-24 h-1.5 bg-primary/20 rounded-full overflow-hidden hidden sm:block">
                                        <div
                                            className="h-full bg-primary transition-all duration-300 ease-out rounded-full"
                                            style={{ width: `${Math.min(progressPercentage, 100)}%` }}
                                        />
                                    </div>
                                    <span className="text-xs text-muted-foreground hidden sm:inline min-w-[60px]">
                                        {getProgressText()}
                                    </span>
                                </div>
                            )}
                            {!isInProgress && (
                                <span className="text-xs text-muted-foreground hidden sm:inline">
                                    {taskStatus}
                                </span>
                            )}
                            {activeWsTaskCount > 1 && (
                                <span className="inline-flex items-center justify-center h-4 min-w-[16px] rounded-full bg-blue-500 text-[10px] font-medium text-white px-1">
                                    {activeWsTaskCount}
                                </span>
                            )}
                        </div>
                    </TooltipTrigger>
                    <TooltipContent side="bottom">
                        <p>{getTooltipContent()}</p>
                        <p className="text-xs text-muted-foreground mt-1">ID: {taskID.slice(0, 8)}...</p>
                    </TooltipContent>
                </Tooltip>
                {isInProgress && (
                    <Tooltip>
                        <TooltipTrigger asChild>
                            <Button
                                variant="ghost"
                                size="icon"
                                onClick={handleCancel}
                                disabled={isCancelling}
                                className="h-6 w-6 text-muted-foreground hover:text-destructive"
                            >
                                <X className="h-3 w-3" />
                            </Button>
                        </TooltipTrigger>
                        <TooltipContent side="bottom">
                            <p>Cancel task</p>
                        </TooltipContent>
                    </Tooltip>
                )}
                <Tooltip>
                    <TooltipTrigger asChild>
                        <Button
                            variant="ghost"
                            size="icon"
                            onClick={handleClearAll}
                            disabled={isClearing}
                            className="h-6 w-6 text-muted-foreground hover:text-destructive"
                        >
                            <Trash2 className="h-3 w-3" />
                        </Button>
                    </TooltipTrigger>
                    <TooltipContent side="bottom">
                        <p>Clear all tasks</p>
                    </TooltipContent>
                </Tooltip>
            </div>
            <TaskProgressDrawer
                open={drawerOpen}
                onOpenChange={setDrawerOpen}
                taskResult={taskResult}
                taskStatus={taskStatus}
                taskID={selectedTaskId ?? taskID}
                onCancel={handleCancel}
                isCancelling={isCancelling}
                wsTasks={wsTasks}
                selectedTaskId={selectedTaskId}
                onSelectTask={setSelectedTaskId}
            />
        </TooltipProvider>
    );
}