backstory/frontend/src/pages/GenerateCandidate.tsx

import React, { useEffect, useState, useRef, useCallback } from 'react';
import Avatar from '@mui/material/Avatar';
import Box from '@mui/material/Box';
import Tooltip from '@mui/material/Tooltip';
import Button from '@mui/material/Button';
import Paper from '@mui/material/Paper';
import IconButton from '@mui/material/IconButton';
import CancelIcon from '@mui/icons-material/Cancel';
import SendIcon from '@mui/icons-material/Send';
import PropagateLoader from 'react-spinners/PropagateLoader';

import { CandidateInfo } from '../components/ui/CandidateInfo';
import { Quote } from 'components/Quote';
import { BackstoryElementProps } from 'components/BackstoryTab';
import { BackstoryTextField, BackstoryTextFieldRef } from 'components/BackstoryTextField';
import { StyledMarkdown } from 'components/StyledMarkdown';
import { Scrollable } from '../components/Scrollable';
import { Pulse } from 'components/Pulse';
import { StreamingResponse } from 'services/api-client';
import {
  ChatMessage,
  ChatMessageUser,
  ChatSession,
  CandidateAI,
  ChatMessageStatus,
  ChatMessageError,
} from 'types/types';
import { useAuth } from 'hooks/AuthContext';
import { Message } from 'components/Message';
import { useAppState } from 'hooks/GlobalContext';

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

const GenerateCandidate = (props: BackstoryElementProps) => {
  const { apiClient, user } = useAuth();
  const { setSnack } = useAppState();
  const [processingMessage, setProcessingMessage] = useState<ChatMessage | null>(null);
  const [processing, setProcessing] = useState<boolean>(false);
  const [generatedUser, setGeneratedUser] = useState<CandidateAI | null>(null);
  const [prompt, setPrompt] = useState<string>('');
  const [resume, setResume] = useState<string | null>(null);
  const [canGenImage, setCanGenImage] = useState<boolean>(false);
  const [timestamp, setTimestamp] = useState<string>('');
  const [shouldGenerateProfile, setShouldGenerateProfile] = useState<boolean>(false);
  const [chatSession, setChatSession] = useState<ChatSession | null>(null);
  const [loading, setLoading] = useState<boolean>(false);

  // Only keep refs that are truly necessary
  const controllerRef = useRef<StreamingResponse>(null);
  const backstoryTextRef = useRef<BackstoryTextFieldRef>(null);

  /* Create the chat session */
  useEffect(() => {
    if (chatSession || loading || !generatedUser) {
      return;
    }

    try {
      setLoading(true);
      apiClient
        .getOrCreateChatSession(
          generatedUser,
          `Profile image generator for ${generatedUser.fullName}`,
          'generate_image'
        )
        .then(session => {
          setChatSession(session);
          setLoading(false);
        });
    } catch (error) {
      setSnack('Unable to load chat session', 'error');
    } finally {
      setLoading(false);
    }
  }, [generatedUser, chatSession, loading, setChatSession, setLoading, setSnack, apiClient]);

  const cancelQuery = useCallback(() => {
    if (controllerRef.current) {
      controllerRef.current.cancel();
      controllerRef.current = null;
      setProcessing(false);
    }
  }, []);

  const onEnter = useCallback(
    (value: string) => {
      if (processing) {
        return;
      }

      const generatePersona = async (prompt: string) => {
        const userMessage: ChatMessageUser = {
          type: 'text',
          role: 'user',
          content: prompt,
          sessionId: '',
          status: 'done',
          timestamp: new Date(),
        };
        setPrompt(prompt || '');
        setProcessing(true);
        setProcessingMessage({
          ...defaultMessage,
          content: 'Generating persona...',
        });
        try {
          const result = await apiClient.createCandidateAI(userMessage);
          console.log(result.message, result);
          setGeneratedUser(result.candidate);
          setResume(result.resume);
          setCanGenImage(true);
          setShouldGenerateProfile(true); // Reset the flag
        } catch (error) {
          console.error(error);
          setPrompt('');
          setResume(null);
          setProcessing(false);
          setProcessingMessage(null);
          setSnack('Unable to generate AI persona', 'error');
        }
      };

      generatePersona(value);
    },
    [processing, apiClient, setSnack]
  );

  const handleSendClick = useCallback(() => {
    const value = (backstoryTextRef.current && backstoryTextRef.current.getAndResetValue()) || '';
    onEnter(value);
  }, [onEnter]);

  // Effect to trigger profile image generation when user data is ready
  useEffect(() => {
    if (!chatSession || !generatedUser?.username) {
      return;
    }
    const username = generatedUser.username;

    if (
      !shouldGenerateProfile ||
      username === '[blank]' ||
      generatedUser?.firstName === '[blank]'
    ) {
      return;
    }

    if (controllerRef.current) {
      console.log('Controller already active, skipping profile generation');
      return;
    }

    setProcessingMessage({
      ...defaultMessage,
      content: 'Starting image generation...',
    });
    setProcessing(true);
    setCanGenImage(false);

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

    controllerRef.current = apiClient.sendMessageStream(chatMessage, {
      onMessage: async (msg: ChatMessage) => {
        console.log(`onMessage: ${msg.type} ${msg.content}`, msg);
        controllerRef.current = null;
        try {
          await apiClient.updateCandidate(generatedUser.id || '', {
            profileImage: 'profile.png',
          });
          const { success, message } = await apiClient.deleteChatSession(chatSession.id || '');
          console.log(
            `Profile generated for ${username} and chat session was ${
              !success ? 'not ' : ''
            } deleted: ${message}}`
          );
          setGeneratedUser({
            ...generatedUser,
            profileImage: 'profile.png',
          } as CandidateAI);
          setCanGenImage(true);
          setShouldGenerateProfile(false);
        } catch (error) {
          console.error(error);
          setSnack(
            `Unable to update ${username} to indicate they have a profile picture.`,
            'error'
          );
        }
      },
      onError: (error: string | ChatMessageError) => {
        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) {
          setSnack(error.content || 'Unknown error generating profile image', 'error');
        } else {
          setSnack(error as string, 'error');
        }
        setProcessingMessage(null);
        setProcessing(false);
        controllerRef.current = null;
        setCanGenImage(true);
        setShouldGenerateProfile(false);
      },
      onComplete: () => {
        setProcessingMessage(null);
        setProcessing(false);
        controllerRef.current = null;
        setCanGenImage(true);
        setShouldGenerateProfile(false);
      },
      onStatus: (status: ChatMessageStatus) => {
        if (status.activity === 'heartbeat' && status.content) {
          setTimestamp(status.timestamp?.toISOString() || '');
        } else if (status.content) {
          setProcessingMessage({ ...defaultMessage, content: status.content });
        }
        console.log(`onStatusChange: ${status}`);
      },
    });
  }, [chatSession, shouldGenerateProfile, generatedUser, prompt, setSnack, apiClient]);

  if (!user?.isAdmin) {
    return <Box>You must be logged in as an admin to generate AI candidates.</Box>;
  }

  return (
    <Box
      className="GenerateCandidate"
      sx={{
        display: 'flex',
        flexDirection: 'column',
        flexGrow: 1,
        gap: 1,
        maxWidth: { xs: '100%', md: '700px', lg: '1024px' },
      }}
    >
      {generatedUser && <CandidateInfo candidate={generatedUser} sx={{ flexShrink: 1 }} />}
      {prompt && <Quote quote={prompt} />}
      {processing && (
        <Box
          sx={{
            display: 'flex',
            flexDirection: 'column',
            alignItems: 'center',
            justifyContent: 'center',
            m: 2,
          }}
        >
          {processingMessage && chatSession && (
            <Message message={processingMessage} {...{ chatSession }} />
          )}
          <PropagateLoader
            size="10px"
            loading={processing}
            aria-label="Loading Spinner"
            data-testid="loader"
          />
        </Box>
      )}
      <Box sx={{ display: 'flex', flexDirection: 'column' }}>
        <Box
          sx={{
            display: 'flex',
            flexDirection: 'row',
            position: 'relative',
          }}
        >
          <Box
            sx={{
              display: 'flex',
              position: 'relative',
              width: 'min-content',
              height: 'min-content',
            }}
          >
            <Avatar
              src={
                generatedUser?.profileImage
                  ? `/api/1.0/candidates/profile/${generatedUser.username}`
                  : ''
              }
              alt={`${generatedUser?.fullName}'s profile`}
              sx={{
                width: 80,
                height: 80,
                border: '2px solid #e0e0e0',
              }}
            />
            {processing && (
              <Pulse
                sx={{
                  position: 'relative',
                  left: '-80px',
                  top: '0px',
                  mr: '-80px',
                }}
                timestamp={timestamp}
              />
            )}
          </Box>

          <Tooltip title={`${generatedUser?.profileImage ? 'Re-' : ''}Generate Picture`}>
            <span style={{ display: 'flex', flexGrow: 1 }}>
              <Button
                sx={{
                  m: 1,
                  gap: 1,
                  justifySelf: 'flex-start',
                  alignSelf: 'center',
                  flexGrow: 0,
                  maxHeight: 'min-content',
                }}
                variant="contained"
                disabled={processing || !canGenImage}
                onClick={() => {
                  setShouldGenerateProfile(true);
                }}
              >
                {generatedUser?.profileImage ? 'Re-' : ''}Generate Picture
                <SendIcon />
              </Button>
            </span>
          </Tooltip>
        </Box>
      </Box>
      {resume && (
        <Paper sx={{ pt: 1, pb: 1, pl: 2, pr: 2 }}>
          <Scrollable sx={{ flexGrow: 1 }}>
            <StyledMarkdown content={resume} />
          </Scrollable>
        </Paper>
      )}
      <BackstoryTextField
        style={{ flexGrow: 0, flexShrink: 1 }}
        ref={backstoryTextRef}
        disabled={processing}
        onEnter={onEnter}
        placeholder='Specify any characteristics you would like the persona to have. For example, "This person likes yo-yos."'
      />
      <Box sx={{ display: 'flex', justifyContent: 'center', flexDirection: 'row' }}>
        <Tooltip title={'Send'}>
          <span style={{ display: 'flex', flexGrow: 1 }}>
            <Button
              sx={{ m: 1, gap: 1, flexGrow: 1 }}
              variant="contained"
              disabled={processing}
              onClick={handleSendClick}
            >
              Generate New Persona
              <SendIcon />
            </Button>
          </span>
        </Tooltip>
        <Tooltip title="Cancel">
          <span style={{ display: 'flex' }}>
            {' '}
            {/* This span is used to wrap the IconButton to ensure Tooltip works even when disabled */}
            <IconButton
              aria-label="cancel"
              onClick={cancelQuery}
              sx={{ display: 'flex', margin: 'auto 0px' }}
              size="large"
              edge="start"
              disabled={controllerRef.current === null || processing === false}
            >
              <CancelIcon />
            </IconButton>
          </span>
        </Tooltip>
      </Box>
      <Box sx={{ display: 'flex', flexGrow: 1 }} />
    </Box>
  );
};

export { GenerateCandidate };