backstory/frontend/src/pages/CandidateChatPage.tsx

import React, { forwardRef, useState, useEffect, useRef } from 'react';
import {
  Box,
  Paper,
  Typography,
  TextField,
  Button,
  List,
  ListItem,
  ListItemText,
  Chip,
  IconButton,
  CircularProgress,
  Divider,
  Card,
  CardContent,
  Avatar,
  Grid
} from '@mui/material';
import { useUser } from 'hooks/useUser';
import { ChatMessageBase, ChatMessage, ChatSession } from 'types/types';
import { ConversationHandle } from 'components/Conversation';
import { BackstoryPageProps } from 'components/BackstoryTab';
import { Message } from 'components/Message';
import { DeleteConfirmation } from 'components/DeleteConfirmation';
import { CandidateSessionsResponse } from 'services/api-client';
import { CandidateInfo } from 'components/CandidateInfo';
import { useNavigate } from 'react-router-dom';

const CandidateChatPage = forwardRef<ConversationHandle, BackstoryPageProps>((props: BackstoryPageProps, ref) => {
  const { apiClient, candidate } = useUser();
  const navigate = useNavigate();
  const {
    setSnack,
    submitQuery,
  } = props;  
  const [sessions, setSessions] = useState<CandidateSessionsResponse | null>(null);
  const [chatSession, setChatSession] = useState<ChatSession | null>(null);
  const [messages, setMessages] = useState([]);
  const [newMessage, setNewMessage] = useState('');
  const [loading, setLoading] = useState(false);
  const [streaming, setStreaming] = useState(false);
  const messagesEndRef = useRef(null);

  // Load sessions for the candidate
  const loadSessions = async () => {
    if (!candidate) return;
    
    try {
      setLoading(true);
      const result = await apiClient.getCandidateChatSessions(candidate.username);
      setSessions(result);
    } catch (error) {
      console.error('Failed to load sessions:', error);
    } finally {
      setLoading(false);
    }
  };

  // Load messages for current session
  const loadMessages = async () => {
    if (!chatSession?.id) return;
    
    try {
      const result = await apiClient.getChatMessages(chatSession.id);
      setMessages(result.data as any);
    } catch (error) {
      console.error('Failed to load messages:', error);
    }
  };

  // Create new session
  const createNewSession = async () => {
    if (!candidate) { return }
    try {
      setLoading(true);
      const newSession = await apiClient.createCandidateChatSession(
        candidate.username,
        'candidate_chat',
        `Interview Discussion - ${candidate.username}`
      );
      setChatSession(newSession);
      setMessages([]);
      await loadSessions(); // Refresh sessions list
    } catch (error) {
      console.error('Failed to create session:', error);
    } finally {
      setLoading(false);
    }
  };

  // Send message
  const sendMessage = async () => {
    if (!newMessage.trim() || !chatSession?.id || streaming) return;

    const messageContent = newMessage;
    setNewMessage('');
    setStreaming(true);

    try {
      await apiClient.sendMessageStream(
        chatSession.id, 
        { prompt: messageContent }, {
        onMessage: (msg) => {
          console.log("onMessage:", msg);
          if (msg.type === "response") {
          setMessages(prev => {
            const filtered = prev.filter((m: any) => m.id !== msg.id);
            return [...filtered, msg].sort((a, b) => 
              new Date(a.timestamp).getTime() - new Date(b.timestamp).getTime()
            ) as any;
          });
          } else {
            console.log(msg);
          }
          },
          onError: (error: string | ChatMessageBase) => {
            console.log("onError:", error);
            setStreaming(false);
          },
          onStreaming: (chunk) => {
            console.log("onStreaming:", chunk);
          },
          onStatusChange: (status) => {
            console.log("onStatusChange:", status);
          },
          onComplete: () => {
            console.log("onComplete");
            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(() => {
    loadSessions();
  }, [candidate]);

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

  if (!candidate) {
    navigate('/find-a-candidate');
  }

  return (
    <Box ref={ref} sx={{ width: "100%", height: "100%", display: "flex", flexGrow: 1, flexDirection: "column" }}>
      {candidate && <CandidateInfo action={`Chat with Backstory about ${candidate.firstName}`} elevation={4} candidate={candidate} sx={{ minHeight: "max-content" }} />}

      < Box sx={{ display: "flex", mt: 1, gap: 1, height: "100%" }}>
        {/* Sessions Sidebar */}
        <Paper sx={{ p: 2, height: '100%', minWidth: { sm: "200px", md: "300px", lg: "400px" }, display: 'flex', flexDirection: 'column' }}>
          <Typography variant="h6" gutterBottom>
            Chat Sessions
            {sessions && (
              <Chip
                size="small"
                label={`${sessions.sessions.total} total`}
                sx={{ ml: 1 }}
              />
            )}
          </Typography>

          <Button
            variant="outlined"
            onClick={createNewSession}
            disabled={loading || !candidate}
            sx={{ mb: 2 }}
          >
            New Session
          </Button>

          <Box sx={{ flexGrow: 1, overflow: 'auto' }}>
            {sessions ? (
              <List>
                {sessions.sessions.data.map((session: any) => (
                  <ListItem
                    key={session.id}
                    // selected={chatSession?.id === session.id}
                    onClick={() => setChatSession(session)}
                    sx={{
                      mb: 1,
                      borderRadius: 1,
                      border: '1px solid',
                      borderColor: chatSession?.id === session.id ? 'primary.main' : 'divider',
                      cursor: 'pointer',
                      '&:hover': {
                        backgroundColor: 'action.hover'
                      }
                    }}
                  >
                    <ListItemText
                      primary={session.title}
                      secondary={`${new Date(session.lastActivity).toLocaleDateString()} • ${session.context.type}`}
                    />
                  </ListItem>
                ))}
              </List>
            ) : (
              <Typography color="text.secondary" align="center">
                Enter a username and click "Load Sessions"
              </Typography>
            )}
          </Box>
        </Paper>

        {/* Chat Interface */}
        <Paper sx={{ display: 'flex', flexDirection: 'column', flexGrow: 1 }}>
          {chatSession?.id ? (
            <>
              {/* Messages Area */}
              <Box sx={{ flexGrow: 1, overflow: 'auto', p: 2 }}>
                {messages.map((message: ChatMessageBase) => (
                  <Message key={message.id} {...{ chatSession, message, setSnack, submitQuery }} />
                ))}

                {streaming && (
                  <Box sx={{ display: 'flex', alignItems: 'center', mb: 2 }}>
                    <Avatar sx={{ mr: 1, bgcolor: 'primary.main' }}>
                      🤖
                    </Avatar>
                    <Card>
                      <CardContent sx={{ display: 'flex', alignItems: 'center', p: 2 }}>
                        <CircularProgress size={16} sx={{ mr: 1 }} />
                        <Typography variant="body2">AI is typing...</Typography>
                      </CardContent>
                    </Card>
                  </Box>
                )}

                <div ref={messagesEndRef} />
              </Box>

              <Divider />

              {/* Message Input */}
              <Box sx={{ p: 2, display: 'flex', gap: 1 }}>
                <TextField
                  fullWidth
                  variant="outlined"
                  placeholder="Type your message about the candidate..."
                  value={newMessage}
                  onChange={(e) => setNewMessage(e.target.value)}
                  onKeyPress={(e) => {
                    if (e.key === 'Enter' && !e.shiftKey) {
                      e.preventDefault();
                      sendMessage();
                    }
                  }}
                  disabled={streaming}
                  multiline
                  maxRows={4}
                />
                <Button
                  variant="contained"
                  onClick={sendMessage}
                  disabled={!newMessage.trim() || streaming}
                  sx={{ minWidth: 'auto', px: 2 }}
                >
                  ▶
                </Button>
              </Box>
            </>
          ) : (
            <Box
              sx={{
                height: '100%',
                display: 'flex',
                alignItems: 'center',
                justifyContent: 'center',
                flexDirection: 'column',
                gap: 2
              }}
            >
              <Typography variant="h1" sx={{ fontSize: 64, color: 'text.secondary' }}>
                🤖
              </Typography>
              <Typography variant="h6" color="text.secondary">
                Select a session to start chatting
              </Typography>
              <Typography variant="body2" color="text.secondary" align="center">
                Create a new session or choose from existing ones to begin discussing the candidate
              </Typography>
            </Box>
          )}
        </Paper>
      </Box>
    </Box>
  );
});

export { CandidateChatPage };