stack/apps/web/src/app/(dashboard)/chat/page.tsx

'use client';

import { useCallback, useEffect, useRef, useState } from 'react';
import { api } from '@/lib/api';
import { destroySocket, getSocket } from '@/lib/socket';
import type { Conversation, Message } from '@/lib/types';
import { ConversationList } from '@/components/chat/conversation-list';
import { MessageBubble } from '@/components/chat/message-bubble';
import { ChatInput } from '@/components/chat/chat-input';
import { StreamingMessage } from '@/components/chat/streaming-message';

export default function ChatPage(): React.ReactElement {
  const [conversations, setConversations] = useState<Conversation[]>([]);
  const [activeId, setActiveId] = useState<string | null>(null);
  const [messages, setMessages] = useState<Message[]>([]);
  const [streamingText, setStreamingText] = useState('');
  const [isStreaming, setIsStreaming] = useState(false);
  const messagesEndRef = useRef<HTMLDivElement>(null);

  // Track the active conversation ID in a ref so socket event handlers always
  // see the current value without needing to be re-registered.
  const activeIdRef = useRef<string | null>(null);
  activeIdRef.current = activeId;

  // Accumulate streamed text in a ref so agent:end can read the full content
  // without stale-closure issues.
  const streamingTextRef = useRef('');

  // Load conversations on mount
  useEffect(() => {
    api<Conversation[]>('/api/conversations')
      .then(setConversations)
      .catch(() => {});
  }, []);

  // Load messages when active conversation changes
  useEffect(() => {
    if (!activeId) {
      setMessages([]);
      return;
    }
    // Clear streaming state when switching conversations
    setIsStreaming(false);
    setStreamingText('');
    streamingTextRef.current = '';
    api<Message[]>(`/api/conversations/${activeId}/messages`)
      .then(setMessages)
      .catch(() => {});
  }, [activeId]);

  // Auto-scroll to bottom
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages, streamingText]);

  // Socket.io setup — connect once for the page lifetime
  useEffect(() => {
    const socket = getSocket();

    function onAgentStart(data: { conversationId: string }): void {
      // Only update state if the event belongs to the currently viewed conversation
      if (activeIdRef.current !== data.conversationId) return;
      setIsStreaming(true);
      setStreamingText('');
      streamingTextRef.current = '';
    }

    function onAgentText(data: { conversationId: string; text: string }): void {
      if (activeIdRef.current !== data.conversationId) return;
      streamingTextRef.current += data.text;
      setStreamingText((prev) => prev + data.text);
    }

    function onAgentEnd(data: { conversationId: string }): void {
      if (activeIdRef.current !== data.conversationId) return;
      const finalText = streamingTextRef.current;
      setIsStreaming(false);
      setStreamingText('');
      streamingTextRef.current = '';
      // Append the completed assistant message to the local message list.
      // The Pi agent session is in-memory so the assistant response is not
      // persisted to the DB — we build the local UI state instead.
      if (finalText) {
        setMessages((prev) => [
          ...prev,
          {
            id: `assistant-${Date.now()}`,
            conversationId: data.conversationId,
            role: 'assistant' as const,
            content: finalText,
            createdAt: new Date().toISOString(),
          },
        ]);
      }
    }

    function onError(data: { error: string; conversationId?: string }): void {
      setIsStreaming(false);
      setStreamingText('');
      streamingTextRef.current = '';
      setMessages((prev) => [
        ...prev,
        {
          id: `error-${Date.now()}`,
          conversationId: data.conversationId ?? '',
          role: 'system' as const,
          content: `Error: ${data.error}`,
          createdAt: new Date().toISOString(),
        },
      ]);
    }

    socket.on('agent:start', onAgentStart);
    socket.on('agent:text', onAgentText);
    socket.on('agent:end', onAgentEnd);
    socket.on('error', onError);

    // Connect if not already connected
    if (!socket.connected) {
      socket.connect();
    }

    return () => {
      socket.off('agent:start', onAgentStart);
      socket.off('agent:text', onAgentText);
      socket.off('agent:end', onAgentEnd);
      socket.off('error', onError);
      // Fully tear down the socket when the chat page unmounts so we get a
      // fresh authenticated connection next time the page is visited.
      destroySocket();
    };
  }, []);

  const handleNewConversation = useCallback(async () => {
    const conv = await api<Conversation>('/api/conversations', {
      method: 'POST',
      body: { title: 'New conversation' },
    });
    setConversations((prev) => [conv, ...prev]);
    setActiveId(conv.id);
    setMessages([]);
  }, []);

  const handleRename = useCallback(async (id: string, title: string) => {
    const updated = await api<Conversation>(`/api/conversations/${id}`, {
      method: 'PATCH',
      body: { title },
    });
    setConversations((prev) => prev.map((c) => (c.id === id ? updated : c)));
  }, []);

  const handleDelete = useCallback(
    async (id: string) => {
      await api<void>(`/api/conversations/${id}`, { method: 'DELETE' });
      setConversations((prev) => prev.filter((c) => c.id !== id));
      if (activeId === id) {
        setActiveId(null);
        setMessages([]);
      }
    },
    [activeId],
  );

  const handleArchive = useCallback(
    async (id: string, archived: boolean) => {
      const updated = await api<Conversation>(`/api/conversations/${id}`, {
        method: 'PATCH',
        body: { archived },
      });
      setConversations((prev) => prev.map((c) => (c.id === id ? updated : c)));
      // If archiving the active conversation, deselect it
      if (archived && activeId === id) {
        setActiveId(null);
        setMessages([]);
      }
    },
    [activeId],
  );

  const handleSend = useCallback(
    async (content: string) => {
      let convId = activeId;

      // Auto-create conversation if none selected
      if (!convId) {
        const autoTitle = content.slice(0, 60);
        const conv = await api<Conversation>('/api/conversations', {
          method: 'POST',
          body: { title: autoTitle },
        });
        setConversations((prev) => [conv, ...prev]);
        setActiveId(conv.id);
        convId = conv.id;
      } else {
        // Auto-title: if the active conversation still has the default "New
        // conversation" title and this is the first message, update the title
        // from the message content.
        const activeConv = conversations.find((c) => c.id === convId);
        if (activeConv?.title === 'New conversation' && messages.length === 0) {
          const autoTitle = content.slice(0, 60);
          api<Conversation>(`/api/conversations/${convId}`, {
            method: 'PATCH',
            body: { title: autoTitle },
          })
            .then((updated) => {
              setConversations((prev) => prev.map((c) => (c.id === convId ? updated : c)));
            })
            .catch(() => {});
        }
      }

      // Optimistic user message in local UI state
      setMessages((prev) => [
        ...prev,
        {
          id: `user-${Date.now()}`,
          conversationId: convId,
          role: 'user' as const,
          content,
          createdAt: new Date().toISOString(),
        },
      ]);

      // Persist the user message to the DB so conversation history is
      // available when the page is reloaded or a new session starts.
      api<Message>(`/api/conversations/${convId}/messages`, {
        method: 'POST',
        body: { role: 'user', content },
      }).catch(() => {
        // Non-fatal: the agent can still process the message even if
        // REST persistence fails.
      });

      // Send to WebSocket — gateway creates/resumes the agent session and
      // streams the response back via agent:start / agent:text / agent:end.
      const socket = getSocket();
      if (!socket.connected) {
        socket.connect();
      }
      socket.emit('message', { conversationId: convId, content });
    },
    [activeId, conversations, messages],
  );

  return (
    <div className="-m-6 flex h-[calc(100vh-3.5rem)]">
      <ConversationList
        conversations={conversations}
        activeId={activeId}
        onSelect={setActiveId}
        onNew={handleNewConversation}
        onRename={handleRename}
        onDelete={handleDelete}
        onArchive={handleArchive}
      />

      <div className="flex flex-1 flex-col">
        {activeId ? (
          <>
            <div className="flex-1 space-y-4 overflow-y-auto p-6">
              {messages.map((msg) => (
                <MessageBubble key={msg.id} message={msg} />
              ))}
              {isStreaming && <StreamingMessage text={streamingText} />}
              <div ref={messagesEndRef} />
            </div>
            <ChatInput onSend={handleSend} disabled={isStreaming} />
          </>
        ) : (
          <div className="flex flex-1 items-center justify-center">
            <div className="text-center">
              <h2 className="text-lg font-medium text-text-secondary">Welcome to Mosaic Chat</h2>
              <p className="mt-1 text-sm text-text-muted">
                Select a conversation or start a new one
              </p>
              <button
                type="button"
                onClick={handleNewConversation}
                className="mt-4 rounded-lg bg-blue-600 px-4 py-2 text-sm font-medium text-white transition-colors hover:bg-blue-700"
              >
                Start new conversation
              </button>
            </div>
          </div>
        )}
      </div>
    </div>
  );
}