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';

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

const GenerateCandidate = (props: BackstoryElementProps) => {
    const { apiClient, user } = useAuth();
    const { setSnack, submitQuery } = props;
    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, submitQuery, setSnack }} />}
                    <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, setSnack, submitQuery }} />
            </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
};