backstory/frontend/src/pages/CandidateChatPage.tsx

import React, { forwardRef, useState, useEffect, useRef } from 'react';
import { Box, Paper, Button, Divider, useTheme, useMediaQuery, Tooltip } from '@mui/material';
import { Send as SendIcon } from '@mui/icons-material';
import { useAuth } from 'hooks/AuthContext';
import {
  ChatMessage,
  ChatSession,
  ChatMessageUser,
  ChatMessageError,
  ChatMessageStreaming,
  ChatMessageStatus,
} from 'types/types';
import { ConversationHandle } from 'components/Conversation';
import { BackstoryPageProps } from 'components/BackstoryTab';
import { Message } from 'components/Message';
import { DeleteConfirmation } from 'components/DeleteConfirmation';
import { CandidateInfo } from 'components/ui/CandidateInfo';
import { useNavigate } from 'react-router-dom';
import { useAppState, useSelectedCandidate } from 'hooks/GlobalContext';
import PropagateLoader from 'react-spinners/PropagateLoader';
import { BackstoryTextField, BackstoryTextFieldRef } from 'components/BackstoryTextField';
import { BackstoryQuery } from 'components/BackstoryQuery';
import { CandidatePicker } from 'components/ui/CandidatePicker';
import { Scrollable } from 'components/Scrollable';

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

const CandidateChatPage = forwardRef<ConversationHandle, BackstoryPageProps>(
  (props: BackstoryPageProps, ref) => {
    const { apiClient } = useAuth();
    const navigate = useNavigate();
    const { selectedCandidate, setSelectedCandidate } = useSelectedCandidate();
    const theme = useTheme();
    const [processingMessage, setProcessingMessage] = useState<
      ChatMessageStatus | ChatMessageError | null
    >(null);
    const [streamingMessage, setStreamingMessage] = useState<ChatMessage | null>(null);
    const backstoryTextRef = useRef<BackstoryTextFieldRef>(null);

    const { setSnack } = useAppState();

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

    // Load messages for current session
    const loadMessages = async () => {
      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);
      }
    };

    const onDelete = async (session: ChatSession) => {
      if (!session.id) {
        return;
      }
      try {
        await apiClient.resetChatSession(session.id);
        // If we're deleting the currently selected session, clear it
        setMessages([]);
        setSnack('Session reset succeeded', 'success');
      } catch (error) {
        console.error('Failed to delete session:', error);
        setSnack('Failed to delete session', 'error');
      }
    };

    // Send message
    const sendMessage = async (message: string) => {
      if (!message.trim() || !chatSession?.id || streaming || !selectedCandidate) 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 ${selectedCandidate.firstName}'s chat session.`,
      });

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

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

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

    // Load sessions when username changes
    useEffect(() => {
      if (!selectedCandidate) return;
      try {
        setLoading(true);
        apiClient
          .getOrCreateChatSession(
            selectedCandidate,
            `Backstory chat with ${selectedCandidate.fullName}`,
            'candidate_chat'
          )
          .then(session => {
            setChatSession(session);
            setLoading(false);
          });
      } catch (error) {
        setSnack('Unable to load chat session', 'error');
      } finally {
        setLoading(false);
      }
    }, [selectedCandidate]);

    // Load messages when session changes
    useEffect(() => {
      if (chatSession?.id) {
        loadMessages();
      }
    }, [chatSession]);

    if (!selectedCandidate) {
      return <CandidatePicker />;
    }

    const welcomeMessage: ChatMessage = {
      sessionId: chatSession?.id || '',
      role: 'information',
      type: 'text',
      status: 'done',
      timestamp: new Date(),
      content: `Welcome to the Backstory Chat about ${selectedCandidate.fullName}. Ask any questions you have about ${selectedCandidate.firstName}.`,
      metadata: null as any,
    };

    return (
      <Box
        ref={ref}
        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',
        }}
      >
        <Paper elevation={2} sx={{ m: 1, p: 1 }}>
          <CandidateInfo
            key={selectedCandidate.username}
            action={`Chat with Backstory about ${selectedCandidate.firstName}`}
            elevation={4}
            candidate={selectedCandidate}
            variant="small"
            sx={{
              flexShrink: 1,
              width: '100%',
              maxHeight: 0,
              minHeight: 'min-content',
            }} // Prevent header from shrinking
          />
          <Button
            sx={{ maxWidth: 'max-content' }}
            onClick={() => {
              setSelectedCandidate(null);
            }}
            variant="contained"
          >
            Change Candidates
          </Button>
        </Paper>
        {/* 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.length === 0 && <Message {...{ chatSession, message: welcomeMessage }} />}
            {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>
        )}
        {selectedCandidate.questions?.length !== 0 &&
          selectedCandidate.questions?.map(q => <BackstoryQuery question={q} />)}
        {/* Fixed Message Input */}
        <Box sx={{ display: 'flex', flexShrink: 1, gap: 1 }}>
          <DeleteConfirmation
            onDelete={() => {
              chatSession && onDelete(chatSession);
            }}
            disabled={!chatSession}
            sx={{ minWidth: 'auto', px: 2, maxHeight: 'min-content' }}
            action="reset"
            label="chat session"
            title="Reset Chat Session"
            message={`Are you sure you want to reset the session? This action cannot be undone.`}
          />
          <BackstoryTextField
            placeholder="Type your message about the candidate..."
            ref={backstoryTextRef}
            onEnter={sendMessage}
            disabled={streaming || loading}
          />
          <Tooltip title="Send">
            <span
              style={{
                minWidth: 'auto',
                maxHeight: 'min-content',
                alignSelf: 'center',
              }}
            >
              <Button
                variant="contained"
                onClick={() => {
                  sendMessage(
                    (backstoryTextRef.current && backstoryTextRef.current.getAndResetValue()) || ''
                  );
                }}
                disabled={streaming || loading}
              >
                <SendIcon />
              </Button>
            </span>
          </Tooltip>
        </Box>
      </Box>
    );
  }
);

export { CandidateChatPage };