veza/apps/web/src/features/chat/hooks/useChat.ts

import { useEffect, useRef, useState, useCallback } from 'react';
import { useAuthStore } from '@/stores/auth';
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 { user } = useAuthStore();
  const userId = user?.id;
  const username = user?.username;
  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,
  };
};