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 { jsonrepair } from 'jsonrepair';

import { CandidateInfo } from '../components/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 { ChatContext, ChatMessage, ChatMessageUser, ChatMessageBase, ChatSession, ChatQuery, Candidate, CandidateAI } from 'types/types';
import { useAuth } from 'hooks/AuthContext';
import { Message } from 'components/Message';
import { Types } from '@uiw/react-json-view';
import { assert } from 'console';

const emptyUser: CandidateAI = {
    userType: "candidate",
    isAI: true,
    description: "[blank]",
    username: "[blank]",
    firstName: "[blank]",
    lastName: "[blank]",
    fullName: "[blank] [blank]",
    questions: [],
    location: {
        city: '[blank]',
        country: '[blank]'
    },
    email: '[blank]',
    createdAt: new Date(),
    updatedAt: new Date(),
    status: "pending",
    skills: [],
    experience: [],
    education: [],
    preferredJobTypes: [],
    languages: [],
    certifications: [],
    isAdmin: false,
    profileImage: undefined,
    ragContentSize: 0
};

const defaultMessage: ChatMessage = {
    type: "preparing", status: "done", sender: "system", sessionId: "", timestamp: new Date(), content: ""
};

const GenerateCandidate = (props: BackstoryElementProps) => {
    const { apiClient, user } = useAuth();
    const { setSnack, submitQuery } = props;
    const [streaming, setStreaming] = useState<boolean>(false);
    const [streamingMessage, setStreamingMessage] = useState<ChatMessage | null>(null);
    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<number>(0);
    const [state, setState] = useState<number>(0); // Replaced stateRef
    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;
        }

        const createChatSession = async () => {
            console.log('Creating chat session');
            try {
                const response: ChatSession = await apiClient.createCandidateChatSession(
                    generatedUser.username,
                    "generate_image",
                    "Profile image generation"
                );
                setChatSession(response);
                console.log(`Chat session created for generate_image`, response);
                setSnack(`Chat session created for generate_image: ${response.id}`);
            } catch (e) {
                console.error(e);
                setSnack("Unable to create image generation session.", "error");
            }
        };

        setLoading(true);
        createChatSession().then(() => { setLoading(false) });
    }, [generatedUser, chatSession, loading, setChatSession, setLoading, setSnack]);

    const generatePersona = async (prompt: string) => {
        const userMessage: ChatMessageUser = {
            content: prompt,
            sessionId: "",
            sender: "user",
            status: "done",
            type: "user",
            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");
        }
    };

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

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

    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);
        setState(3);

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

        controllerRef.current = apiClient.sendMessageStream(chatMessage, {
            onMessage: async (msg: ChatMessage) => {
                console.log(`onMessage: ${msg.type} ${msg.content}`, msg);
                if (msg.type === "heartbeat") {
                    const heartbeat = JSON.parse(msg.content);
                    setTimestamp(heartbeat.timestamp);
                }
                if (msg.type === "thinking") {
                    const status = JSON.parse(msg.content);
                    setProcessingMessage({ ...defaultMessage, content: status.message });
                }
                if (msg.type === "response") {
                    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);
                        setState(0);
                        setCanGenImage(true);
                        setShouldGenerateProfile(false);
                    } catch (error) {
                        console.error(error);
                        setSnack(`Unable to update ${username} to indicate they have a profile picture.`, "error");
                    }
                }
            },
            onError: (error) => {
                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);
                setStreaming(false);
                setProcessing(false);
                controllerRef.current = null;
                setState(0);
                setCanGenImage(true);
                setShouldGenerateProfile(false);
            },
            onComplete: () => {
                setProcessingMessage(null);
                setStreaming(false);
                setProcessing(false);
                controllerRef.current = null;
                setState(0);
                setCanGenImage(true);
                setShouldGenerateProfile(false);
            },
            onStatusChange: (status: string) => {
                console.log(`onStatusChange: ${status}`);
            },
        });
    }, [chatSession, shouldGenerateProfile, generatedUser, prompt, setSnack]);

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