"use client";

import { useEffect, useRef, useCallback, useState } from "react";
import { WS_BASE_URL } from "@/lib/config";
import type { WSMessage, WSMessageType } from "@/types";

interface UseWebSocketOptions {
  onMessage?: (message: WSMessage) => void;
  onConnect?: () => void;
  onDisconnect?: () => void;
  onError?: (error: Event) => void;
  autoConnect?: boolean;
  authToken?: string | null;
}

export function useWebSocket(options: UseWebSocketOptions = {}) {
  const { onMessage, onConnect, onDisconnect, onError, autoConnect = true, authToken = null } = options;
  
  const wsRef = useRef<WebSocket | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [clientId, setClientId] = useState<string | null>(null);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const reconnectAttemptsRef = useRef(0);
  const maxReconnectAttempts = 5;

  const connect = useCallback(() => {
    const openConnection = () => {
      if (wsRef.current?.readyState === WebSocket.OPEN) return;

      const id = crypto.randomUUID();
      setClientId(id);

      const wsUrl = authToken
        ? `${WS_BASE_URL}/${id}?token=${encodeURIComponent(authToken)}`
        : `${WS_BASE_URL}/${id}`;

      const ws = new WebSocket(wsUrl);
      wsRef.current = ws;

      ws.onopen = () => {
        setIsConnected(true);
        reconnectAttemptsRef.current = 0;
        onConnect?.();
      };

      ws.onclose = () => {
        setIsConnected(false);
        onDisconnect?.();

        // Attempt reconnection
        if (reconnectAttemptsRef.current < maxReconnectAttempts) {
          const delay = Math.min(1000 * Math.pow(2, reconnectAttemptsRef.current), 30000);
          reconnectTimeoutRef.current = setTimeout(() => {
            reconnectAttemptsRef.current++;
            openConnection();
          }, delay);
        }
      };

      ws.onerror = (error) => {
        onError?.(error);
      };

      ws.onmessage = (event) => {
        try {
          const message = JSON.parse(event.data) as WSMessage;
          onMessage?.(message);
        } catch (e) {
          console.error("Failed to parse WebSocket message:", e);
        }
      };
    };

    openConnection();
  }, [authToken, onConnect, onDisconnect, onError, onMessage]);

  const disconnect = useCallback(() => {
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
    }
    reconnectAttemptsRef.current = maxReconnectAttempts; // Prevent reconnection
    wsRef.current?.close();
    wsRef.current = null;
    setIsConnected(false);
  }, []);

  const sendMessage = useCallback((type: WSMessageType, taskId?: string, data: Record<string, unknown> = {}) => {
    if (wsRef.current?.readyState !== WebSocket.OPEN) {
      console.error("WebSocket is not connected");
      return false;
    }

    const message: WSMessage = {
      type,
      task_id: taskId,
      data,
      timestamp: new Date().toISOString(),
    };

    wsRef.current.send(JSON.stringify(message));
    return true;
  }, []);

  useEffect(() => {
    if (autoConnect) {
      connect();
    }

    return () => {
      disconnect();
    };
  }, [autoConnect, connect, disconnect]);

  return {
    isConnected,
    clientId,
    connect,
    disconnect,
    sendMessage,
  };
}
