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 { Candidate } from '../types/types';
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 } from 'types/types';
import { useAuth } from 'hooks/AuthContext';

const emptyUser: Candidate = {
    userType: "candidate",
    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: []
};

const GenerateCandidate = (props: BackstoryElementProps) => {
    const { apiClient } = useAuth();
    const { setSnack, submitQuery } = props;
    const [streaming, setStreaming] = useState<string>('');
    const [processing, setProcessing] = useState<boolean>(false);
    const [user, setUser] = useState<Candidate | null>(null);
    const [prompt, setPrompt] = useState<string>('');
    const [resume, setResume] = useState<string>('');
    const [canGenImage, setCanGenImage] = useState<boolean>(false);
    const [status, setStatus] = useState<string>('');
    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);

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

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

        const createChatSession = async () => {
            try {
                const chatContext: ChatContext = { type: "generate_persona" };
                const response: ChatSession = await apiClient.createChatSession(chatContext);
                setChatSession(response);
                setSnack(`Chat session created for generate_persona: ${response.id}`);
            } catch (e) {
                console.error(e);
                setSnack("Unable to create chat session.", "error");
            }
        };

        createChatSession();
    }, [chatSession, setChatSession]);

    const generatePersona = useCallback((query: ChatQuery) => {
        if (!chatSession || !chatSession.id) {
            return;
        }
        const sessionId: string = chatSession.id;

        setPrompt(query.prompt || '');
        setState(0);
        setStatus("Generating persona...");
        setUser(emptyUser);
        setStreaming('');
        setResume('');
        setProcessing(true);
        setCanGenImage(false);
        setShouldGenerateProfile(false); // Reset the flag

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

        const streamResponse = apiClient.sendMessageStream(chatMessage, {
            onMessage: (chatMessage: ChatMessage) => {
                console.log('Message:', chatMessage);
                // Update UI with partial content
            },
            onStatusChange: (status) => {
                console.log('Status changed:', status);
                // Update UI status indicator
            },
            onComplete: () => {
                console.log('Content complete');
            },
            onWarn: (warning) => {
                console.log("Warning:", warning);
            },
            onError: (error: string | ChatMessageBase) => {
                // Type-guard to determine if this is a ChatMessageBase or a string
                if (typeof error === "object" && error !== null && "content" in error) {
                    console.log("Error message:", error);
                } else {
                    console.log("Error string:", error);
                }
            },
            onStreaming: (chunk) => {
                console.log("Streaming: ", chunk);
            }
        });
        // controllerRef.current = streamQueryResponse({
        //     query,
        //     type: "persona",
        //     connectionBase,
        //     onComplete: (msg) => {
        //         switch (msg.status) {
        //             case "partial":
        //             case "done":
        //                 setState(currentState => {
        //                     switch (currentState) {
        //                         case 0: /* Generating persona */
        //                             let partialUser = JSON.parse(jsonrepair((msg.response || '').trim()));
        //                             if (!partialUser.fullName) {
        //                                 partialUser.fullName = `${partialUser.firstName} ${partialUser.lastName}`;
        //                             }
        //                             console.log("Setting final user data:", partialUser);
        //                             setUser({ ...partialUser });
        //                             return 1; /* Generating resume */
        //                         case 1: /* Generating resume */
        //                             setResume(msg.response || '');
        //                             return 2; /* RAG generation */
        //                         case 2: /* RAG generation */
        //                             return 3; /* Image generation */
        //                         default:
        //                             return currentState;
        //                     }
        //                 });

        //                 if (msg.status === "done") {
        //                     setProcessing(false);
        //                     setCanGenImage(true);
        //                     setStatus('');
        //                     controllerRef.current = null;
        //                     setState(0);
        //                     // Set flag to trigger profile generation after user state updates
        //                     console.log("Persona generation complete, setting shouldGenerateProfile flag");
        //                     setShouldGenerateProfile(true);
        //                 }
        //                 break;
        //             case "thinking":
        //                 setStatus(msg.response || '');
        //                 break;

        //             case "error":
        //                 console.log(`Error generating persona: ${msg.response}`);
        //                 setSnack(msg.response || "", "error");
        //                 setProcessing(false);
        //                 setUser(emptyUser);
        //                 controllerRef.current = null;
        //                 setState(0);
        //                 break;
        //         }
        //     },
        //     onStreaming: (chunk) => {
        //         setStreaming(chunk);
        //     }
        // });
    }, [setSnack]);

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

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

    const handleSendClick = useCallback(() => {
        const value = (backstoryTextRef.current && backstoryTextRef.current.getAndResetValue()) || "";
        generatePersona({ prompt: value });
    }, [generatePersona]);

    // Effect to trigger profile generation when user data is ready
    useEffect(() => {
        console.log("useEffect triggered - shouldGenerateProfile:", shouldGenerateProfile, "user:", user?.username, user?.firstName);
        if (shouldGenerateProfile && user?.username !== "[blank]" && user?.firstName !== "[blank]") {
            console.log("Triggering profile generation with updated user data:", user);
            if (controllerRef.current) {
                console.log("Controller already active, skipping profile generation");
                return;
            }

            // Don't generate if we still have blank user data
            if (user?.username === "[blank]" || user?.firstName === "[blank]") {
                console.log("Cannot generate profile: user data not ready");
                return;
            }

            const imagePrompt = `A photorealistic profile picture of a ${user?.age} year old ${user?.gender?.toLocaleLowerCase()} ${user?.ethnicity?.toLocaleLowerCase()} person. ${prompt}`
            setStatus('Starting image generation...');
            setProcessing(true);
            setCanGenImage(false);
            setState(3);
            const start = Date.now();

            // controllerRef.current = streamQueryResponse({
            //     query: {
            //         prompt: imagePrompt,
            //         agentOptions: {
            //             username: user?.username,
            //             filename: "profile.png"
            //         }
            //     },
            //     type: "image",
            //     sessionId,
            //     connectionBase,
            //     onComplete: (msg) => {
            //         // console.log("Profile generation response:", msg);
            //         switch (msg.status) {
            //             case "partial":
            //             case "done":
            //                 if (msg.status === "done") {
            //                     setProcessing(false);
            //                     controllerRef.current = null;
            //                     setState(0);
            //                     setCanGenImage(true);
            //                     setShouldGenerateProfile(false);
            //                     setUser({
            //                       ...(user ? user : emptyUser),
            //                         hasProfile: true
            //                     });
            //                 }
            //                 break;
            //             case "error":
            //                 console.log(`Error generating profile: ${msg.response} after ${Date.now() - start}`);
            //                 setSnack(msg.response || "", "error");
            //                 setProcessing(false);
            //                 controllerRef.current = null;
            //                 setState(0);
            //                 setCanGenImage(true);
            //                 setShouldGenerateProfile(false);
            //                 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;
            //         }
            //     }
            // });
        }
    }, [shouldGenerateProfile, user, prompt, setSnack]);

    // Handle streaming updates based on current state
    useEffect(() => {
        if (streaming.trim().length === 0) {
            return;
        }

        try {
            switch (state) {
                case 0: /* Generating persona */
                    const partialUser = {...emptyUser, ...JSON.parse(jsonrepair(`${streaming.trim()}...`))};
                    if (!partialUser.fullName) {
                        partialUser.fullName = `${partialUser.firstName} ${partialUser.lastName}`;
                    }
                    setUser(partialUser);
                    break;
                case 1: /* Generating resume */
                    setResume(streaming);
                    break;
                case 3: /* RAG streaming */
                    break;
                case 4: /* Image streaming */
                    break;
            }
        } catch {
            // Ignore JSON parsing errors during streaming
        }
    }, [streaming, state]);

    return (
      <Box className="GenerateCandidate" sx={{ 
        display: "flex", 
        flexDirection: "column", 
        flexGrow: 1, 
        gap: 1,
        maxWidth: { xs: '100%', md: '700px', lg: '1024px' },
        }}>
        {user && <CandidateInfo 
                candidate={user} 
          sx={{flexShrink: 1}}/>
        }
        { prompt && 
          <Quote quote={prompt}/>
        }
        {processing && 
          <Box sx={{
              display: "flex",
              flexDirection: "column",
              alignItems: "center",
              justifyContent: "center",
              m: 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}
                  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={user?.hasProfile ? `/api/u/${user.username}/profile` : ''}
                            alt={`${user?.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={`${user?.hasProfile ? '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); }}>
                                {user?.hasProfile ? '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
};