import { useState, useRef, useCallback, useEffect } from 'react';
import { 
  WorkflowEvent, 
  WorkflowDag, 
  TaskStatus, 
  StartWorkflowRequest, 
  StartWorkflowResponse 
} from '@/types/workflow';

export type WorkflowConnectionStatus = 'idle' | 'connecting' | 'connected' | 'disconnected' | 'error';

interface UseWorkflowReturn {
  // State
  status: WorkflowConnectionStatus;
  requestId: string | null;
  dag: WorkflowDag | null;
  taskStates: Record<string, TaskStatus>;
  taskOutputs: Record<string, string>; // Accumulates streaming content
  error: string | null;
  finalResult: any | null;

  // Actions
  // Returns StartWorkflowResponse to allow caller to handle redirects (e.g. symbol normalization)
  startWorkflow: (params: StartWorkflowRequest) => Promise<StartWorkflowResponse | undefined>;
  connectToWorkflow: (requestId: string) => void;
  disconnect: () => void;
}

export function useWorkflow(): UseWorkflowReturn {
  const [status, setStatus] = useState<WorkflowConnectionStatus>('idle');
  const [requestId, setRequestId] = useState<string | null>(null);
  const [dag, setDag] = useState<WorkflowDag | null>(null);
  const [taskStates, setTaskStates] = useState<Record<string, TaskStatus>>({});
  const [taskOutputs, setTaskOutputs] = useState<Record<string, string>>({});
  const [error, setError] = useState<string | null>(null);
  const [finalResult, setFinalResult] = useState<any | null>(null);

  // Ref for EventSource to handle cleanup
  const eventSourceRef = useRef<EventSource | null>(null);
  
  // Refs for state that updates frequently to avoid closure staleness in event handlers if needed
  // (Though in this React pattern, simple state updates usually suffice unless high freq)

  const disconnect = useCallback(() => {
    if (eventSourceRef.current) {
      eventSourceRef.current.close();
      eventSourceRef.current = null;
    }
    setStatus('disconnected');
  }, []);

  const handleEvent = useCallback((eventData: WorkflowEvent) => {
    switch (eventData.type) {
      case 'WorkflowStarted':
        setDag(eventData.payload.task_graph);
        // Initialize states based on graph
        const initialStates: Record<string, TaskStatus> = {};
        eventData.payload.task_graph.nodes.forEach(node => {
          initialStates[node.id] = node.initial_status;
        });
        setTaskStates(initialStates);
        break;

      case 'TaskStateChanged':
        setTaskStates(prev => ({
          ...prev,
          [eventData.payload.task_id]: eventData.payload.status
        }));
        break;

      case 'TaskStreamUpdate':
        setTaskOutputs(prev => ({
          ...prev,
          [eventData.payload.task_id]: (prev[eventData.payload.task_id] || '') + eventData.payload.content_delta
        }));
        break;

      case 'WorkflowStateSnapshot':
        // Restore full state
        setDag(eventData.payload.task_graph);
        setTaskStates(eventData.payload.tasks_status);
        // Restore outputs if present
        const outputs: Record<string, string> = {};
        Object.entries(eventData.payload.tasks_output).forEach(([k, v]) => {
          if (v) outputs[k] = v;
        });
        setTaskOutputs(prev => ({ ...prev, ...outputs }));
        break;

      case 'WorkflowCompleted':
        setFinalResult(eventData.payload.result_summary);
        disconnect(); // Close connection on completion
        break;

      case 'WorkflowFailed':
        setError(eventData.payload.reason);
        // We might want to keep connected or disconnect depending on if retry is possible
        // For now, treat fatal error as disconnect reason
        if (eventData.payload.is_fatal) {
          disconnect();
          setStatus('error');
        }
        break;
    }
  }, [disconnect]);

  const connectToWorkflow = useCallback((id: string) => {
    if (eventSourceRef.current) {
      eventSourceRef.current.close();
    }

    setRequestId(id);
    setStatus('connecting');
    setError(null);

    try {
      const es = new EventSource(`/api/workflow/events/${id}`);
      eventSourceRef.current = es;

      es.onopen = () => {
        setStatus('connected');
      };

      es.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data) as WorkflowEvent;
          handleEvent(data);
        } catch (e) {
          console.error('Failed to parse workflow event:', e);
        }
      };

      es.onerror = (e) => {
        console.error('Workflow SSE error:', e);
        // EventSource automatically retries, but we might want to handle it explicitly
        // For now, let's assume if readyState is CLOSED, it's a fatal error
        if (es.readyState === EventSource.CLOSED) {
           setStatus('error');
           setError('Connection lost');
           es.close();
        }
      };

    } catch (e) {
      console.error('Failed to create EventSource:', e);
      setStatus('error');
      setError(e instanceof Error ? e.message : 'Connection initialization failed');
    }
  }, [handleEvent]);

  const startWorkflow = useCallback(async (params: StartWorkflowRequest) => {
    setStatus('connecting');
    setError(null);
    setDag(null);
    setTaskStates({});
    setTaskOutputs({});
    setFinalResult(null);

    try {
      const res = await fetch('/api/workflow/start', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(params),
      });

      if (!res.ok) {
        const errorBody = await res.json().catch(() => ({}));
        throw new Error(errorBody.error || `HTTP ${res.status}`);
      }

      const data: StartWorkflowResponse = await res.json();
      
      // Start listening
      connectToWorkflow(data.request_id);
      
      return data; // Return response so UI can handle symbol normalization redirection
      
    } catch (e) {
      setStatus('error');
      setError(e instanceof Error ? e.message : 'Failed to start workflow');
      return undefined;
    }
  }, [connectToWorkflow]);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      if (eventSourceRef.current) {
        eventSourceRef.current.close();
      }
    };
  }, []);

  return {
    status,
    requestId,
    dag,
    taskStates,
    taskOutputs,
    error,
    finalResult,
    startWorkflow,
    connectToWorkflow,
    disconnect
  };
}