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 {
  ConversationSidebar,
  type ConversationSidebarRef,
} from '@/components/chat/conversation-sidebar';
import { MessageBubble } from '@/components/chat/message-bubble';
import { ChatInput } from '@/components/chat/chat-input';
import { StreamingMessage } from '@/components/chat/streaming-message';

interface ModelInfo {
  id: string;
  provider: string;
  name: string;
  reasoning: boolean;
  contextWindow: number;
  maxTokens: number;
  inputTypes: ('text' | 'image')[];
  cost: { input: number; output: number; cacheRead: number; cacheWrite: number };
}

interface ProviderInfo {
  id: string;
  name: string;
  available: boolean;
  models: ModelInfo[];
}

export default function ChatPage(): React.ReactElement {
  const [activeId, setActiveId] = useState<string | null>(null);
  const [messages, setMessages] = useState<Message[]>([]);
  const [streamingText, setStreamingText] = useState('');
  const [isStreaming, setIsStreaming] = useState(false);
  const [isSidebarOpen, setIsSidebarOpen] = useState(true);
  const [models, setModels] = useState<ModelInfo[]>([]);
  const [selectedModelId, setSelectedModelId] = useState('');
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const sidebarRef = useRef<ConversationSidebarRef>(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('');

  useEffect(() => {
    const savedState = window.localStorage.getItem('mosaic-sidebar-open');
    if (savedState !== null) {
      setIsSidebarOpen(savedState === 'true');
    }
  }, []);

  useEffect(() => {
    window.localStorage.setItem('mosaic-sidebar-open', String(isSidebarOpen));
  }, [isSidebarOpen]);

  useEffect(() => {
    api<ProviderInfo[]>('/api/providers')
      .then((providers) => {
        const availableModels = providers
          .filter((provider) => provider.available)
          .flatMap((provider) => provider.models);
        setModels(availableModels);
        setSelectedModelId((current) => current || availableModels[0]?.id || '');
      })
      .catch(() => {
        setModels([]);
        setSelectedModelId('');
      });
  }, []);

  // 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(),
          },
        ]);
        sidebarRef.current?.refresh();
      }
    }

    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 (projectId?: string | null) => {
    const conv = await api<Conversation>('/api/conversations', {
      method: 'POST',
      body: { title: 'New conversation', projectId: projectId ?? null },
    });

    sidebarRef.current?.addConversation({
      id: conv.id,
      title: conv.title,
      projectId: conv.projectId,
      updatedAt: conv.updatedAt,
      archived: conv.archived,
    });

    setActiveId(conv.id);
    setMessages([]);
    setIsSidebarOpen(true);
  }, []);

  const handleSend = useCallback(
    async (content: string, options?: { modelId?: 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 },
        });
        sidebarRef.current?.addConversation({
          id: conv.id,
          title: conv.title,
          projectId: conv.projectId,
          updatedAt: conv.updatedAt,
          archived: conv.archived,
        });
        setActiveId(conv.id);
        convId = conv.id;
      } else if (messages.length === 0) {
        // Auto-title the initial placeholder conversation from the first user message.
        const autoTitle = content.slice(0, 60);
        api<Conversation>(`/api/conversations/${convId}`, {
          method: 'PATCH',
          body: { title: autoTitle },
        })
          .then(() => sidebarRef.current?.refresh())
          .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,
        modelId: (options?.modelId ?? selectedModelId) || undefined,
      });
    },
    [activeId, messages, selectedModelId],
  );

  return (
    <div
      className="-m-6 flex h-[calc(100vh-3.5rem)] overflow-hidden"
      style={{ background: 'var(--bg-deep, var(--color-surface-bg, #0a0f1a))' }}
    >
      <ConversationSidebar
        ref={sidebarRef}
        isOpen={isSidebarOpen}
        onClose={() => setIsSidebarOpen(false)}
        currentConversationId={activeId}
        onSelectConversation={(conversationId) => {
          setActiveId(conversationId);
          setMessages([]);
          if (conversationId && window.innerWidth < 768) {
            setIsSidebarOpen(false);
          }
        }}
        onNewConversation={(projectId) => {
          void handleNewConversation(projectId);
        }}
      />

      <div className="flex min-w-0 flex-1 flex-col">
        <div
          className="flex items-center gap-3 border-b px-4 py-3"
          style={{ borderColor: 'var(--border)' }}
        >
          <button
            type="button"
            onClick={() => setIsSidebarOpen((open) => !open)}
            className="rounded-lg border p-2 transition-colors"
            style={{
              borderColor: 'var(--border)',
              background: 'var(--surface)',
              color: 'var(--text)',
            }}
            aria-label={isSidebarOpen ? 'Close conversation sidebar' : 'Open conversation sidebar'}
          >
            <svg viewBox="0 0 24 24" className="h-4 w-4" fill="none" stroke="currentColor">
              <path strokeWidth="2" strokeLinecap="round" d="M4 7h16M4 12h16M4 17h16" />
            </svg>
          </button>
          <div>
            <h1 className="text-sm font-semibold" style={{ color: 'var(--text)' }}>
              Mosaic Chat
            </h1>
            <p className="text-xs" style={{ color: 'var(--muted)' }}>
              {activeId ? 'Active conversation selected' : 'Choose or start a conversation'}
            </p>
          </div>
        </div>

        {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}
              isStreaming={isStreaming}
              models={models}
              selectedModelId={selectedModelId}
              onModelChange={setSelectedModelId}
            />
          </>
        ) : (
          <div className="flex flex-1 items-center justify-center px-6">
            <div
              className="max-w-md rounded-2xl border px-8 py-10 text-center"
              style={{
                borderColor: 'var(--border)',
                background: 'var(--surface)',
              }}
            >
              <h2 className="text-lg font-medium" style={{ color: 'var(--text)' }}>
                Welcome to Mosaic Chat
              </h2>
              <p className="mt-1 text-sm" style={{ color: 'var(--muted)' }}>
                Select a conversation or start a new one
              </p>
              <button
                type="button"
                onClick={() => {
                  void handleNewConversation();
                }}
                className="mt-4 rounded-lg px-4 py-2 text-sm font-medium text-white transition-colors"
                style={{ background: 'var(--primary)' }}
              >
                Start new conversation
              </button>
            </div>
          </div>
        )}
      </div>
    </div>
  );
}