import { useState, useEffect, useRef, useCallback } from 'react';

export interface PromptProgressState {
  percent: number;
  processed: number;
  total: number;
  cache: number;
  timeMs: number;
}

export interface ToolCallStreamEntry {
  name: string;
  rawArgs: string;
}

export interface JobStreamState {
  promptProgress: PromptProgressState | null;
  streamingText: string;
  toolCallStream: Record<string, ToolCallStreamEntry>;
  connected: boolean;
}

export function useJobStream(taskId: number | null, jobStatus: string | null | undefined): JobStreamState {
  const [promptProgress, setPromptProgress] = useState<PromptProgressState | null>(null);
  const [streamingText, setStreamingText] = useState('');
  const [toolCallStream, setToolCallStream] = useState<Record<string, ToolCallStreamEntry>>({});
  const [connected, setConnected] = useState(false);
  const esRef = useRef<EventSource | null>(null);

  const isActive = jobStatus === 'running' || jobStatus === 'dispatching';

  const cleanup = useCallback(() => {
    if (esRef.current) {
      esRef.current.close();
      esRef.current = null;
    }
    setConnected(false);
  }, []);

  useEffect(() => {
    if (!taskId || !isActive) {
      cleanup();
      setPromptProgress(null);
      setStreamingText('');
      setToolCallStream({});
      return;
    }

    const es = new EventSource(`/api/local/tasks/${taskId}/stream`);
    esRef.current = es;

    es.onopen = () => setConnected(true);

    es.onmessage = (e) => {
      try {
        const data = JSON.parse(e.data);
        switch (data.type) {
          case 'prompt_progress':
            setPromptProgress({
              percent: data.percent,
              processed: data.processed,
              total: data.total,
              cache: data.cache ?? 0,
              timeMs: data.timeMs ?? 0,
            });
            break;
          case 'text_delta':
            setPromptProgress(null);
            setStreamingText(prev => prev + (data.text ?? ''));
            break;
          case 'tool_use_delta':
            setPromptProgress(null);
            setToolCallStream(prev => {
              const callId = data.callId ?? '';
              const existing = prev[callId];
              // chunk is a full snapshot of args-so-far — REPLACE, not append.
              // This is what lets a late-attaching stream still get the prefix.
              return {
                ...prev,
                [callId]: {
                  name: data.name ?? existing?.name ?? '',
                  rawArgs: data.chunk ?? existing?.rawArgs ?? '',
                },
              };
            });
            break;
          case 'tool_use':
            setStreamingText('');
            setPromptProgress(null);
            setToolCallStream(prev => {
              if (!data.callId || !(data.callId in prev)) return prev;
              const next = { ...prev };
              delete next[data.callId];
              return next;
            });
            break;
          case 'done':
            setStreamingText('');
            setPromptProgress(null);
            setToolCallStream({});
            cleanup();
            break;
        }
      } catch { /* ignore malformed */ }
    };

    es.onerror = () => {
      setConnected(false);
      // EventSource auto-reconnects
    };

    return cleanup;
  }, [taskId, isActive, cleanup]);

  return { promptProgress, streamingText, toolCallStream, connected };
}
