import { useEffect, useRef, useState, useCallback } from 'react';
import { useAuthStore } from '@/features/auth/store/authStore'; // Our main auth store
import { useChatStore } from '../store/chatStore';
import { apiClient } from '@/lib/apiClient';
import { OutgoingMessage, IncomingMessage } from '../types'; 
import { v4 as uuidv4 } from 'uuid'; // For message IDs

export const useChat = () => {
  const { userId, username } = useAuthStore();
  const {
    wsToken,
    wsUrl,
    wsStatus,
    setWsStatus,
    addMessage,
    currentConversationId,
    loadMessages,
  } = useChatStore();
  
  const ws = useRef<WebSocket | null>(null);
  const [messagesToSend, setMessagesToSend] = useState<OutgoingMessage[]>([]); // Queue for messages to send

  const connect = useCallback(() => {
    if (!wsToken || !wsUrl || ws.current?.readyState === WebSocket.OPEN) return;

    setWsStatus('connecting');
    const fullWsUrl = `${wsUrl}?token=${wsToken}`; // Assuming WS server is at root of wsUrl
    ws.current = new WebSocket(fullWsUrl);

    ws.current.onopen = () => {
      setWsStatus('connected');
      console.log('WebSocket connected');
      // Send any queued messages
      setMessagesToSend((prev) => {
        prev.forEach((msg) => ws.current?.send(JSON.stringify(msg)));
        return [];
      });
    };

    ws.current.onmessage = (event) => {
      const data = JSON.parse(event.data);
      if (data.type === 'NewMessage') {
        const message: IncomingMessage = data; // Cast to IncomingMessage
        if (message.conversation_id === currentConversationId) {
          addMessage({
            id: message.message_id,
            conversation_id: message.conversation_id,
            sender_id: message.sender_id,
            sender_username: message.sender_username || 'Unknown', // Need sender_username from backend
            content: message.content,
            created_at: message.created_at,
          });
        }
      }
      // Handle other incoming message types (ActionConfirmed, Error, Pong)
    };

    ws.current.onclose = () => {
      setWsStatus('disconnected');
      console.log('WebSocket disconnected');
      // Optional: Reconnect logic
    };

    ws.current.onerror = (error) => {
      setWsStatus('error');
      console.error('WebSocket error:', error);
      ws.current?.close();
    };

  }, [wsToken, wsUrl, setWsStatus, addMessage, currentConversationId]);

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

  useEffect(() => {
    if (wsToken && wsUrl && wsStatus === 'disconnected') {
      connect();
    }
    // Clean up on unmount
    return () => {
      disconnect();
    };
  }, [wsToken, wsUrl, wsStatus, connect, disconnect]);

  const sendMessage = useCallback((content: string) => {
    if (!ws.current || ws.current.readyState !== WebSocket.OPEN || !currentConversationId || !userId) {
      console.warn('WebSocket not open, cannot send message');
      // Queue message to send later
      setMessagesToSend((prev) => [...prev, {
        type: 'SendMessage',
        conversation_id: currentConversationId || uuidv4(), // Fallback
        content,
        parent_message_id: null,
      }]);
      return;
    }

    const message: OutgoingMessage = {
      type: 'SendMessage',
      conversation_id: currentConversationId,
      content,
      parent_message_id: null,
    };
    ws.current.send(JSON.stringify(message));
  }, [ws.current, currentConversationId, userId]);

  // TODO: Add fetchHistory function
  const fetchHistory = useCallback(async (conversationId: string) => {
    try {
      const response = await apiClient.get(`/conversations/${conversationId}/history`);
      loadMessages(conversationId, response.data.messages);
    } catch (error) {
      console.error('Failed to fetch chat history:', error);
    }
  }, [loadMessages]);

  return {
    wsStatus,
    connect,
    disconnect,
    sendMessage,
    fetchHistory,
  };
};