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

type WebSocketStatus = 'connecting' | 'connected' | 'disconnected' | 'error';

interface WebSocketMessage {
  status: 'pending' | 'processing' | 'completed' | 'failed';
  stage?: string;
  progress?: number;
  message?: string;
  audio_url?: string;
  error?: string;
}

export function useGenerationWebSocket(generationId: string, isActive: boolean) {
  const [status, setStatus] = useState<WebSocketStatus>('disconnected');
  const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null);
  const wsRef = useRef<WebSocket | null>(null);

  useEffect(() => {
    // Only connect if the generation is active (pending/processing)
    if (!isActive) {
      return;
    }

    // Prevent duplicate connections
    if (wsRef.current && (wsRef.current.readyState === WebSocket.OPEN || wsRef.current.readyState === WebSocket.CONNECTING)) {
      return;
    }

    const connect = () => {
      try {
        setStatus('connecting');
        const apiBase = process.env.NEXT_PUBLIC_API_URL || 'http://localhost:8000';
        // Convert http(s) to ws(s)
        const wsBase = apiBase.replace(/^http/, 'ws');
        const wsUrl = `${wsBase}/api/v1/ws/generations/${generationId}`;

        const ws = new WebSocket(wsUrl);

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

        ws.onmessage = (event) => {
          try {
            const data = JSON.parse(event.data);
            setLastMessage(data);
          } catch {
            // Invalid JSON - ignore
          }
        };

        ws.onclose = () => {
          setStatus('disconnected');
        };

        ws.onerror = () => {
          setStatus('error');
        };

        wsRef.current = ws;
      } catch {
        setStatus('error');
      }
    };

    connect();

    return () => {
      const ws = wsRef.current;
      wsRef.current = null;
      if (ws) {
        ws.close();
      }
    };
  }, [generationId, isActive]);

  return { status, lastMessage };
}