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 DelegateStreamEntry {
  delegateRunId: string;
  parentRunId: string | null;
  depth: number;
  description: string;
  status: 'running' | 'success' | 'aborted' | 'needs_user_input';
  text: string;
  currentTool: string | null;
  /** SSE delegate_lifecycle の originJobId フィールド。サブタスクグループ分類に使用。 */
  originJobId?: string | null;
}

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

function emptyEntry(delegateRunId: string): DelegateStreamEntry {
  return {
    delegateRunId, parentRunId: null, depth: 0, description: '',
    status: 'running', text: '', currentTool: null, originJobId: null,
  };
}

/** SSE 1 イベントを delegateStreams マップへ畳み込む純関数。 */
export function reduceDelegateStreams(
  prev: Record<string, DelegateStreamEntry>,
  data: { type: string; delegateRunId?: string; parentRunId?: string | null;
    depth?: number; description?: string; status?: DelegateStreamEntry['status'];
    text?: string; toolName?: string; originJobId?: string | null },
): Record<string, DelegateStreamEntry> {
  const id = data.delegateRunId;
  if (!id) return prev;
  const existing = prev[id] ?? emptyEntry(id);
  switch (data.type) {
    case 'delegate_lifecycle':
      return { ...prev, [id]: {
        ...existing,
        parentRunId: data.parentRunId ?? existing.parentRunId,
        depth: data.depth ?? existing.depth,
        description: data.description || existing.description,
        status: data.status ?? existing.status,
        originJobId: data.originJobId !== undefined ? data.originJobId : existing.originJobId,
      } };
    case 'delegate_text_delta':
      return { ...prev, [id]: { ...existing, text: existing.text + (data.text ?? ''), currentTool: null } };
    case 'delegate_tool':
      return { ...prev, [id]: { ...existing, currentTool: data.toolName || null } };
    default:
      return prev;
  }
}

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 [delegateStreams, setDelegateStreams] = useState<Record<string, DelegateStreamEntry>>({});
  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({});
      setDelegateStreams({});
      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 'delegate_lifecycle':
          case 'delegate_text_delta':
          case 'delegate_tool':
            setDelegateStreams(prev => reduceDelegateStreams(prev, data));
            break;
          case 'done':
            setStreamingText('');
            setPromptProgress(null);
            setToolCallStream({});
            setDelegateStreams({});
            cleanup();
            break;
        }
      } catch { /* ignore malformed */ }
    };

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

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

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