DreamChat/apps/frontend/src/pages/Chat.tsx

import { useEffect, useState, useRef } from 'react';
import { useParams, useNavigate, Link } from 'react-router-dom';
import { useChatStore } from '../stores/chatStore';
import { useCharacterStore } from '../stores/characterStore';
import { useAuthStore } from '../stores/authStore';
import type { Message } from '../types';
import { io, Socket } from 'socket.io-client';
import { chatControllerSendMessage } from '../api/generated/conversations/conversations';

const WS_URL = (import.meta.env as unknown as ImportMetaEnv).VITE_WS_URL || 'http://localhost:3000';

export function Chat() {
  const { characterId, conversationId } = useParams<{ characterId?: string; conversationId?: string }>();
  const navigate = useNavigate();
  const { logout } = useAuthStore();
  const { currentCharacter, getCharacter } = useCharacterStore();
  const { 
    currentConversation, 
    isStreaming,
    error, 
    createConversation, 
    getConversation,
    addMessage,
    setStreaming,
    clearError 
  } = useChatStore();
  
  const [message, setMessage] = useState('');
  const [streamingContent, setStreamingContent] = useState('');
  const [socketConnected, setSocketConnected] = useState(false);
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const socketRef = useRef<Socket | null>(null);

  // Initialize socket connection
  useEffect(() => {
    const token = localStorage.getItem('accessToken');
    if (!token) return;

    const socket = io(`${WS_URL}/chat`, {
      auth: { token: `Bearer ${token}` },
    });

    socketRef.current = socket;

    socket.on('connect', () => {
      console.log('Connected to chat server');
      setSocketConnected(true);
    });

    socket.on('disconnect', () => {
      console.log('Disconnected from chat server');
      setSocketConnected(false);
    });

    socket.on('connect_error', (err) => {
      console.error('Socket connection error:', err.message);
      setSocketConnected(false);
    });

    socket.on('message_chunk', (data: { conversationId: string; chunk: { content: string } }) => {
      if (data.conversationId === conversationId) {
        setStreamingContent((prev) => prev + data.chunk.content);
      }
    });

    socket.on('message_complete', () => {
      setStreaming(false);
      setStreamingContent('');
      if (conversationId) {
        getConversation(conversationId);
      }
    });

    socket.on('message', (data: { message: { assistantMessage?: Message } }) => {
      if (data.message.assistantMessage) {
        addMessage(data.message.assistantMessage);
      }
    });

    socket.on('error', (data: { message: string }) => {
      console.error('Socket error:', data.message);
      setStreaming(false);
    });

    return () => {
      socket.disconnect();
    };
  }, [conversationId, addMessage, getConversation, setStreaming]);

  // Join conversation room
  useEffect(() => {
    if (socketRef.current && conversationId) {
      socketRef.current.emit('join_conversation', { conversationId });
      return () => {
        socketRef.current?.emit('leave_conversation', { conversationId });
      };
    }
  }, [conversationId]);

  // Load character and conversation
  useEffect(() => {
    if (characterId) {
      getCharacter(characterId);
      // Create new conversation
      createConversation({ characterId }).then((conv) => {
        navigate(`/conversations/${conv.id}`, { replace: true });
      });
    } else if (conversationId) {
      getConversation(conversationId);
    }
  }, [characterId, conversationId, getCharacter, createConversation, getConversation, navigate]);

  // Scroll to bottom when messages change
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [currentConversation?.messages, streamingContent]);

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!message.trim() || !conversationId || isStreaming) return;

    const content = message.trim();
    setMessage('');
    setStreamingContent('');
    setStreaming(true);

    // Check if socket is connected, otherwise fallback to HTTP
    if (socketRef.current?.connected) {
      socketRef.current.emit('send_message', {
        conversationId,
        content,
      });
    } else {
      // Fallback to HTTP API when socket is not connected
      try {
        const result = await chatControllerSendMessage(conversationId, { content });
        
        // Add messages to the conversation
        if (result.userMessage) {
          addMessage(result.userMessage as Message);
        }
        if (result.assistantMessage) {
          addMessage(result.assistantMessage as Message);
        }
        setStreaming(false);
      } catch (error) {
        console.error('Failed to send message:', error);
        setStreaming(false);
      }
    }
  };

  const handleLogout = () => {
    logout();
    navigate('/login');
  };

  const formatTime = (dateString: string) => {
    return new Date(dateString).toLocaleTimeString('en-US', {
      hour: '2-digit',
      minute: '2-digit',
    });
  };

  const messages = currentConversation?.messages || [];

  return (
    <div className="h-screen flex flex-col bg-gray-50">
      <header className="bg-white shadow flex-shrink-0">
        <div className="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 py-4 flex justify-between items-center">
          <div className="flex items-center space-x-4">
            <Link to="/conversations" className="text-gray-600 hover:text-gray-900">
              ← Back
            </Link>
            <div>
              <h1 className="text-lg font-semibold text-gray-900">
                {currentConversation?.title || 'Chat'}
              </h1>
              {currentCharacter && (
                <p className="text-sm text-gray-500">
                  with {currentCharacter.name}
                </p>
              )}
            </div>
          </div>
          <button
            onClick={handleLogout}
            className="text-sm text-gray-600 hover:text-gray-900"
          >
            Logout
          </button>
        </div>
      </header>

      <main className="flex-1 overflow-hidden flex flex-col max-w-7xl mx-auto w-full px-4 sm:px-6 lg:px-8 py-4">
        {error && (
          <div className="mb-4 bg-red-50 text-red-700 p-3 rounded flex-shrink-0">
            {error}
            <button onClick={clearError} className="ml-2 text-sm underline">Dismiss</button>
          </div>
        )}

        <div className="flex-1 overflow-y-auto bg-white rounded-lg shadow p-4 space-y-4">
          {messages.length === 0 && !isStreaming && (
            <div className="text-center text-gray-500 py-8">
              <p>Start a conversation with {currentCharacter?.name || 'your character'}</p>
            </div>
          )}

          {messages.map((msg) => (
            <div
              key={msg.id}
              className={`flex ${msg.role === 'user' ? 'justify-end' : 'justify-start'}`}
            >
              <div
                className={`max-w-[70%] rounded-lg px-4 py-2 ${
                  msg.role === 'user'
                    ? 'bg-indigo-600 text-white'
                    : 'bg-gray-100 text-gray-900'
                }`}
              >
                <p className="whitespace-pre-wrap">{msg.content}</p>
                <span className={`text-xs ${msg.role === 'user' ? 'text-indigo-200' : 'text-gray-500'} block mt-1`}>
                  {formatTime(msg.createdAt)}
                </span>
              </div>
            </div>
          ))}

          {isStreaming && streamingContent && (
            <div className="flex justify-start">
              <div className="max-w-[70%] rounded-lg px-4 py-2 bg-gray-100 text-gray-900">
                <p className="whitespace-pre-wrap">{streamingContent}</p>
                <span className="text-xs text-gray-500 block mt-1">typing...</span>
              </div>
            </div>
          )}

          <div ref={messagesEndRef} />
        </div>

        <form onSubmit={handleSubmit} className="mt-4 flex-shrink-0 flex space-x-4">
          <input
            type="text"
            value={message}
            onChange={(e) => setMessage(e.target.value)}
            placeholder="Type your message..."
            disabled={isStreaming}
            className="flex-1 px-4 py-2 border border-gray-300 rounded-lg focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:border-transparent disabled:opacity-50"
          />
          <button
            type="submit"
            disabled={!message.trim() || isStreaming}
            className="px-6 py-2 bg-indigo-600 text-white rounded-lg hover:bg-indigo-700 focus:outline-none focus:ring-2 focus:ring-indigo-500 disabled:opacity-50 disabled:cursor-not-allowed"
          >
            {isStreaming ? 'Sending...' : 'Send'}
          </button>
        </form>
      </main>
    </div>
  );
}