persona-community-2/apps/creator-ui/src/pages/ChatPage.tsx

import { useRef, useEffect, useState, useCallback, useMemo } from 'react';
import { useAuth } from '@persona-community-2/auth';
import { useChat } from '@persona-community-2/realtime';
import {
  Card,
  CardHeader,
  CardTitle,
  CardDescription,
  CardContent,
  ChatBubble,
  ChatInput,
  Badge,
  ProviderBadge,
} from '@persona-community-2/ui';

interface TimelineMessage {
  id: string;
  content: string;
  role: 'user' | 'assistant' | 'system';
  timestamp: Date;
  provider?: string;
  isStreaming?: boolean;
}

export function ChatPage() {
  const { user, getToken } = useAuth();
  const messagesEndRef = useRef<HTMLDivElement>(null);

  // API base URL from environment
  const apiBaseUrl = import.meta.env.VITE_API_URL || '';

  const authHeaders = useMemo(() => {
    const token = getToken();
    return token ? { Authorization: `Bearer ${token}` } : undefined;
  }, [getToken]);

  const {
    messages,
    aiMessages,
    streamingMessages,
    sendMessage,
    connectionState,
    onlineUsers,
  } = useChat({
    endpoint: `${apiBaseUrl}/api/{{SERVICE_NAME}}/chat/messages`,
    sseEndpoint: `${apiBaseUrl}/api/{{SERVICE_NAME}}/events`,
    channel: 'channel:general',
    userId: user?.id || 'anonymous',
    userName: user?.name || user?.email || 'Anonymous',
    headers: authHeaders,
  });

  // Track send errors for user feedback
  const [sendError, setSendError] = useState<string | null>(null);

  // Merge user messages + AI messages into a single sorted timeline
  const timeline = useMemo<TimelineMessage[]>(() => {
    const combined: TimelineMessage[] = [];

    for (const msg of messages) {
      combined.push({
        id: msg.id,
        content: msg.content,
        role: msg.userId === user?.id ? 'user' : 'assistant',
        timestamp: new Date(msg.timestamp),
      });
    }

    for (const msg of aiMessages) {
      combined.push({
        id: msg.id,
        content: msg.content,
        role: 'assistant',
        timestamp: new Date(msg.timestamp),
        provider: msg.provider,
      });
    }

    // Add in-progress streaming messages
    for (const [, stream] of streamingMessages) {
      combined.push({
        id: stream.streamId,
        content: stream.content,
        role: 'assistant',
        timestamp: new Date(stream.timestamp),
        isStreaming: true,
      });
    }

    combined.sort((a, b) => a.timestamp.getTime() - b.timestamp.getTime());
    return combined;
  }, [messages, aiMessages, streamingMessages, user?.id]);

  // Handle sending a message (wraps async sendMessage for ChatInput)
  const handleSendMessage = useCallback((content: string) => {
    sendMessage(content).catch(() => {
      setSendError('Failed to send message. Please try again.');
      setTimeout(() => setSendError(null), 3000);
    });
  }, [sendMessage]);

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

  const connectionBadge = () => {
    switch (connectionState) {
      case 'connected':
        return <Badge variant="success">Connected</Badge>;
      case 'connecting':
        return <Badge variant="warning">Connecting...</Badge>;
      case 'disconnected':
        return <Badge variant="error">Disconnected</Badge>;
      case 'error':
        return <Badge variant="error">Error</Badge>;
      default:
        return null;
    }
  };

  return (
    <div className="h-[calc(100vh-8rem)] flex flex-col">
      <Card className="flex-1 flex flex-col">
        <CardHeader className="flex-shrink-0 flex flex-row items-center justify-between">
          <div>
            <CardTitle>AI Chat</CardTitle>
            <CardDescription>
              Chat with AI in real-time
            </CardDescription>
          </div>
          <div className="flex items-center gap-3">
            <span className="text-sm text-[var(--text-muted)]">
              {onlineUsers.length} online
            </span>
            {connectionBadge()}
          </div>
        </CardHeader>

        <CardContent className="flex-1 flex flex-col min-h-0">
          {/* Messages area */}
          <div className="flex-1 overflow-y-auto space-y-4 pr-2 mb-4">
            {timeline.length === 0 ? (
              <div className="h-full flex items-center justify-center">
                <p className="text-[var(--text-muted)] text-sm">
                  No messages yet. Start the conversation!
                </p>
              </div>
            ) : (
              timeline.map((msg) => (
                <div key={msg.id}>
                  <ChatBubble
                    role={msg.role}
                    content={msg.content}
                    timestamp={msg.timestamp}
                    isStreaming={msg.isStreaming}
                  />
                  {msg.provider && (
                    <div className={msg.role === 'user' ? 'text-right' : 'text-left'}>
                      <ProviderBadge provider={msg.provider} className="mt-1" />
                    </div>
                  )}
                </div>
              ))
            )}
            <div ref={messagesEndRef} />
          </div>

          {/* Input area */}
          <div className="flex-shrink-0 space-y-2">
            {sendError && (
              <div className="px-3 py-2 text-sm text-[var(--error)] bg-[var(--error)]/10 rounded-lg">
                {sendError}
              </div>
            )}
            <ChatInput
              onSubmit={handleSendMessage}
              disabled={connectionState !== 'connected'}
              placeholder={
                connectionState === 'connected'
                  ? 'Type a message... (Cmd+Enter to send)'
                  : 'Connecting...'
              }
            />
          </div>
        </CardContent>
      </Card>
    </div>
  );
}