import { useEffect, useRef } from 'react';
import { useCommStore } from '@/store/commStore';
import type { Conversation, Message } from '@/store/commStore';

/**
 * useCommWebSocket — manages a WebSocket connection to /api/v1/comm/ws.
 *
 * Handles real-time events: new messages, conversation updates, typing
 * indicators, and streaming tokens. Auto-reconnects with exponential backoff.
 */
export function useCommWebSocket() {
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeout = useRef<number | undefined>(undefined);
  const addMessage = useCommStore((s) => s.addMessage);
  const updateConversation = useCommStore((s) => s.updateConversation);
  const setTyping = useCommStore((s) => s.setTyping);
  const activeConversationId = useCommStore((s) => s.activeConversationId);
  const setMessages = useCommStore((s) => s.setMessages);

  useEffect(() => {
    let reconnectAttempts = 0;
    const maxReconnectDelay = 30000;

    function connect() {
      const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
      const ws = new WebSocket(`${protocol}//${window.location.host}/api/v1/comm/ws`);
      wsRef.current = ws;

      ws.onopen = () => {
        reconnectAttempts = 0;
        console.log('Comm WebSocket connected');
      };

      ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          switch (data.type) {
            case 'message.new': {
              if (data.message && data.conversation_id) {
                addMessage(data.conversation_id, data.message as Message);
              }
              break;
            }
            case 'conversation.updated': {
              if (data.conversation) {
                updateConversation(data.conversation as Conversation);
              }
              break;
            }
            case 'message.streaming': {
              // Update last message content with streaming token
              if (data.conversation_id && data.token) {
                const msgs = useCommStore.getState().messages[data.conversation_id] || [];
                if (msgs.length > 0) {
                  const lastMsg = msgs[msgs.length - 1];
                  const updatedMsg = {
                    ...lastMsg,
                    content: lastMsg.content + data.token,
                  };
                  setMessages(data.conversation_id, [...msgs.slice(0, -1), updatedMsg]);
                }
              }
              break;
            }
            case 'message.streaming.done': {
              // Replace streaming content with final message
              if (data.conversation_id && data.message) {
                const msgs = useCommStore.getState().messages[data.conversation_id] || [];
                if (msgs.length > 0) {
                  setMessages(data.conversation_id, [...msgs.slice(0, -1), data.message as Message]);
                }
              }
              break;
            }
            case 'typing': {
              if (data.conversation_id) {
                const currentTyping = useCommStore.getState().typingUsers[data.conversation_id] || [];
                if (data.is_typing) {
                  if (!currentTyping.includes(data.user_id)) {
                    setTyping(data.conversation_id, [...currentTyping, data.user_id]);
                  }
                } else {
                  setTyping(
                    data.conversation_id,
                    currentTyping.filter((uid) => uid !== data.user_id),
                  );
                }
              }
              break;
            }
            case 'pong':
              break;
            default:
              break;
          }
        } catch (err) {
          console.error('Comm WebSocket message parse error:', err);
        }
      };

      ws.onclose = () => {
        console.log('Comm WebSocket disconnected');
        const delay = Math.min(1000 * Math.pow(2, reconnectAttempts), maxReconnectDelay);
        reconnectAttempts++;
        reconnectTimeout.current = window.setTimeout(connect, delay);
      };

      ws.onerror = (error) => {
        console.error('Comm WebSocket error:', error);
      };
    }

    connect();

    // Ping interval to keep connection alive
    const pingInterval = setInterval(() => {
      if (wsRef.current?.readyState === WebSocket.OPEN) {
        wsRef.current.send(JSON.stringify({ type: 'ping' }));
      }
    }, 30000);

    return () => {
      clearInterval(pingInterval);
      if (reconnectTimeout.current) clearTimeout(reconnectTimeout.current);
      wsRef.current?.close();
    };
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []);

  // Subscribe to active conversation when it changes
  useEffect(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN && activeConversationId) {
      wsRef.current.send(
        JSON.stringify({ type: 'subscribe', conversation_id: activeConversationId }),
      );
    }
  }, [activeConversationId]);

  return wsRef;
}
