backstory/frontend/src/components/ui/TransientChat.tsx

import React, { forwardRef, useState, useEffect, useRef, JSX, useImperativeHandle } from 'react';
import { Box } from '@mui/material';
import { useAuth } from 'hooks/AuthContext';
import {
  ChatMessage,
  ChatSession,
  ChatMessageUser,
  ChatMessageError,
  ChatMessageStreaming,
  ChatMessageStatus,
  ChatMessageMetaData,
  Candidate,
  ChatQuery,
} from 'types/types';
import { ConversationHandle } from 'components/Conversation';
import { Message } from 'components/Message';
import { useAppState } from 'hooks/GlobalContext';
import PropagateLoader from 'react-spinners/PropagateLoader';
import { Scrollable } from 'components/Scrollable';

const emptyMetadata: ChatMessageMetaData = {
  model: 'qwen2.5',
  temperature: 0,
  maxTokens: 0,
  topP: 0,
  frequencyPenalty: 0,
  presencePenalty: 0,
  stopSequences: [],
  usage: {
    evalCount: 0,
    evalDuration: 0,
    promptEvalCount: 0,
    promptEvalDuration: 0,
  },
};

const defaultMessage: ChatMessage = {
  status: 'done',
  type: 'text',
  sessionId: '',
  timestamp: new Date(),
  content: '',
  role: 'user',
  metadata: emptyMetadata,
};

interface TransientChatProps {
  id: string;
}

const TransientChat = forwardRef<ConversationHandle, TransientChatProps>(
  (props, ref): JSX.Element => {
    const { id } = props;
    const { apiClient, user } = useAuth();
    const [processingMessage, setProcessingMessage] = useState<
      ChatMessageStatus | ChatMessageError | null
    >(null);
    const [streamingMessage, setStreamingMessage] = useState<ChatMessage | null>(null);

    const { setSnack } = useAppState();

    const [chatSession, setChatSession] = useState<ChatSession | null>(null);
    const [messages, setMessages] = useState<ChatMessage[]>([]);
    const [streaming, setStreaming] = useState<boolean>(false);
    const messagesEndRef = useRef<HTMLDivElement>(null);

    useImperativeHandle(ref, () => ({
      submitQuery: (query: ChatQuery): void => {
        sendMessage(query.prompt);
      },
      fetchHistory: (): void => {
        console.log('fetchHistory called, but not implemented');
      },
    }));

    const chatHandlers = {
      onMessage: (msg: ChatMessage): void => {
        setMessages(prev => {
          const filtered = prev.filter(m => m.id !== msg.id);
          return [...filtered, msg] as ChatMessage[];
        });
        setStreamingMessage(null);
        setProcessingMessage(null);
      },
      onError: (error: string | ChatMessageError): void => {
        console.log('onError:', error);
        // Type-guard to determine if this is a ChatMessageBase or a string
        if (typeof error === 'object' && error !== null && 'content' in error) {
          setProcessingMessage(error);
        } else {
          setProcessingMessage({
            ...defaultMessage,
            status: 'error',
            content: error,
          });
        }
        setStreaming(false);
      },
      onStreaming: (chunk: ChatMessageStreaming): void => {
        // console.log("onStreaming:", chunk);
        setStreamingMessage({
          ...chunk,
          role: 'assistant',
          metadata: emptyMetadata,
        });
      },
      onStatus: (status: ChatMessageStatus): void => {
        setProcessingMessage(status);
      },
      onComplete: (): void => {
        console.log('onComplete');
        setStreamingMessage(null);
        setProcessingMessage(null);
        setStreaming(false);
      },
    };

    // Send message
    const sendMessage = async (message: string): Promise<void> => {
      if (!message.trim() || !chatSession?.id || streaming) return;

      const messageContent = message;
      setStreaming(true);

      const chatMessage: ChatMessageUser = {
        sessionId: chatSession.id,
        role: 'user',
        content: messageContent,
        status: 'done',
        type: 'text',
        timestamp: new Date(),
      };

      setProcessingMessage({
        ...defaultMessage,
        status: 'status',
        activity: 'info',
        content: `Establishing connection with chat session.`,
      });

      setMessages(prev => {
        const filtered = prev.filter(m => m.id !== chatMessage.id);
        return [...filtered, chatMessage] as ChatMessage[];
      });

      try {
        apiClient.sendMessageStream(chatMessage, chatHandlers);
      } catch (error) {
        console.error('Failed to send message:', error);
        setStreaming(false);
      }
    };

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

    // Load sessions when username changes
    useEffect(() => {
      if (!user) return;
      try {
        apiClient
          .getOrCreateChatSession(user as Candidate, `Transient chat - ${id}`, 'candidate_chat')
          .then(session => {
            setChatSession(session);
          });
      } catch (error) {
        setSnack('Unable to load chat session', 'error');
      }
    }, [user, apiClient, setSnack]);

    // Load messages when session changes
    useEffect(() => {
      const loadMessages = async (): Promise<void> => {
        if (!chatSession?.id) return;

        try {
          const result = await apiClient.getChatMessages(chatSession.id);
          const chatMessages: ChatMessage[] = result.data;
          setMessages(chatMessages);
          setProcessingMessage(null);
          setStreamingMessage(null);
          console.log(`getChatMessages returned ${chatMessages.length} messages.`, chatMessages);
        } catch (error) {
          console.error('Failed to load messages:', error);
        }
      };

      if (chatSession?.id) {
        loadMessages();
      }
    }, [chatSession, apiClient]);

    return (
      <Box
        sx={{
          display: 'flex',
          flexDirection: 'column',
          height: '100%' /* Restrict to main-container's height */,
          width: '100%',
          minHeight: 0 /* Prevent flex overflow */,
          maxHeight: 'min-content',
          '& > *:not(.Scrollable)': {
            flexShrink: 0 /* Prevent shrinking */,
          },
          position: 'relative',
        }}
      >
        {/* Chat Interface */}
        {/* Scrollable Messages Area */}
        {chatSession && (
          <Scrollable
            sx={{
              position: 'relative',
              maxHeight: '100%',
              width: '100%',
              display: 'flex',
              flexGrow: 1,
              flex: 1 /* Take remaining space in some-container */,
              overflowY: 'auto' /* Scroll if content overflows */,
              pt: 2,
              pl: 1,
              pr: 1,
              pb: 2,
            }}
          >
            {messages.map((message: ChatMessage) => (
              <Message key={message.id} {...{ chatSession, message }} />
            ))}
            {processingMessage !== null && (
              <Message {...{ chatSession, message: processingMessage }} />
            )}
            {streamingMessage !== null && (
              <Message {...{ chatSession, message: streamingMessage }} />
            )}
            {streaming && (
              <Box
                sx={{
                  display: 'flex',
                  flexDirection: 'column',
                  alignItems: 'center',
                  justifyContent: 'center',
                  m: 1,
                }}
              >
                <PropagateLoader
                  size="10px"
                  loading={streaming}
                  aria-label="Loading Spinner"
                  data-testid="loader"
                />
              </Box>
            )}
            <div ref={messagesEndRef} />
          </Scrollable>
        )}
      </Box>
    );
  }
);
TransientChat.displayName = 'TransientChat';
export { TransientChat };