backstory/frontend/src/components/GenerateImage.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 './CandidateInfo';
import { Query } from '../types/types'
import { Quote } from 'components/Quote';
import { streamQueryResponse, StreamQueryController } from './streamQueryResponse';
import { connectionBase } from 'Global';
import { User } from '../types/types';
import { BackstoryElementProps } from 'components/BackstoryTab';
import { BackstoryTextField, BackstoryTextFieldRef } from 'components/BackstoryTextField';
import { jsonrepair } from 'jsonrepair';
import { StyledMarkdown } from 'components/StyledMarkdown';
import { Scrollable } from './Scrollable';
import { Pulse } from 'components/Pulse';
import { useUser } from './UserContext';

interface GenerateImageProps extends BackstoryElementProps {
    prompt: string
};

const GenerateImage = (props: GenerateImageProps) => {
    const { user } = useUser();
    const {sessionId, setSnack, prompt} = props;
    const [processing, setProcessing] = useState<boolean>(false);
    const [status, setStatus] = useState<string>('');
    const [timestamp, setTimestamp] = useState<number>(0);
    const [image, setImage] = useState<string>('');

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

    // Effect to trigger profile generation when user data is ready
    useEffect(() => {
        if (controllerRef.current) {
            console.log("Controller already active, skipping profile generation");
            return;
        }
        if (!prompt) {
            return;
        }
        setStatus('Starting image generation...');
        setProcessing(true);
        const start = Date.now();

        controllerRef.current = streamQueryResponse({
            query: {
                prompt: prompt,
                agentOptions: {
                    username: user?.username,
                }
            },
            type: "image",
            sessionId,
            connectionBase,
            onComplete: (msg) => {
                switch (msg.status) {
                    case "partial":
                    case "done":
                        if (msg.status === "done") {
                            if (!msg.response) {
                                setSnack("Image generation failed", "error");
                            } else {
                                setImage(msg.response);
                            }
                            setProcessing(false);
                            controllerRef.current = null;
                        }
                        break;
                    case "error":
                        console.log(`Error generating profile: ${msg.response} after ${Date.now() - start}`);
                        setSnack(msg.response || "", "error");
                        setProcessing(false);
                        controllerRef.current = null;
                        break;
                    default:
                        let data: any = {};
                        try {
                            data = typeof msg.response === 'string' ? JSON.parse(msg.response) : msg.response;
                        } catch (e) {
                            data = { message: msg.response };
                        }
                        if (msg.status !== "heartbeat") {
                            console.log(data);
                        }
                        if (data.timestamp) {
                            setTimestamp(data.timestamp);
                        } else {
                            setTimestamp(Date.now())
                        }
                        if (data.message) {
                            setStatus(data.message);
                        }
                        break;
                }
            }
        });
    }, [user, prompt, sessionId, setSnack]);

    if (!sessionId) {
        return <></>;
    }

    return (
      <Box className="GenerateImage" sx={{ 
        display: "flex", 
        flexDirection: "column", 
        flexGrow: 1, 
        gap: 1,
        maxWidth: { xs: '100%', md: '700px', lg: '1024px' },
        minHeight: "max-content",
        }}>
            {image !== '' && <img alt={prompt} src={`${image}/${sessionId}`} />}
        { prompt && 
          <Quote size={processing ? "normal" : "small"} quote={prompt} sx={{ "& *": { color: "#2E2E2E !important" }}}/>
        }
        {processing && 
          <Box sx={{
              display: "flex",
              flexDirection: "column",
              alignItems: "center",
              justifyContent: "center",
              m: 0,
              gap: 1,
              minHeight: "min-content",
              mb: 2
          }}>
              { status && 
              <Box sx={{ display: "flex", flexDirection: "column"}}>
                <Box sx={{ fontSize: "0.5rem"}}>Generation status</Box>
                <Box sx={{ fontWeight: "bold"}}>{status}</Box>
              </Box>
              }
              <PropagateLoader
                  size="10px"
                  loading={processing}
                  color="white"
                  aria-label="Loading Spinner"
                  data-testid="loader"
              />
          </Box>
        }
        </Box>);
};

export {
  GenerateImage
};