import { createContext, ReactNode, useCallback, useContext, useEffect, useRef, useState } from "react";
import { WebSocketContext } from "./WebSocketContext";
import { AccountContext } from "./AccountContext";
import { ConversationContext } from "./ConversationContext";
import { ConversationPrompt, ConversationPromptType } from "../data";
import { axiosInstance } from "../../axiosApi";
import { AxiosResponse } from "axios";
import { AnalyticsEvents, trackEvent } from "../utils/analytics";
import {
    ActivityHistoryEntry,
    parseCitationsData,
    parseStatusData,
    parseVersionedFrame,
} from "../utils/wsFrames";

type MessageProviderProps ={
    children? : ReactNode;
}

type IMessageContext = {
    stateMessage: string;
    setStateMessage: (message: string) => void;
    conversationDetails: ConversationPrompt[];
    setConversationDetails: (conversationPrompts: ConversationPrompt[]) => void;
    isGeneratingMessage: boolean;
    /** True when a stream was interrupted by a socket drop; cleared on refetch/retry. */
    streamInterrupted: boolean;
    clearStreamInterrupted: () => void;
    /** Live activity status from versioned "status" WS frames (#96). */
    activityStage: string | null;
    activityLabel: string | null;
    activityDetail: string | null;
    activityHistory: ActivityHistoryEntry[];
    clearActivity: () => void;
}

const initialValues = {
    stateMessage: '',
    setStateMessage: () => {},
    conversationDetails: [],
    setConversationDetails: () => {},
    isGeneratingMessage: false,
    streamInterrupted: false,
    clearStreamInterrupted: () => {},
    activityStage: null,
    activityLabel: null,
    activityDetail: null,
    activityHistory: [],
    clearActivity: () => {},
}

const MessageContext = createContext<IMessageContext>(initialValues);

const STREAM_INTERRUPTED_MESSAGE =
    '_Response interrupted — connection lost. Re-send your message if needed._';

const MessageProvider = ( {children}: MessageProviderProps) => {
    const [subscribe, unsubscribe, , , , , reconnectGeneration] = useContext(WebSocketContext)
    const { account } = useContext(AccountContext)
    const {selectedConversation, setSelectedConversation} = useContext(ConversationContext);

    const [stateMessage, setStateMessage] = useState<string>('')
    const [conversationDetails, setConversationDetails] = useState<ConversationPrompt[]>([])
    const [isGeneratingMessage, setIsGeneratingMessage] = useState<boolean>(false)
    const [streamInterrupted, setStreamInterrupted] = useState<boolean>(false)
    const [activityStage, setActivityStage] = useState<string | null>(null)
    const [activityLabel, setActivityLabel] = useState<string | null>(null)
    const [activityDetail, setActivityDetail] = useState<string | null>(null)
    const [activityHistory, setActivityHistory] = useState<ActivityHistoryEntry[]>([])

    const messageRef = useRef('')
    const messageResponsePart = useRef(0);
    const conversationRef = useRef(conversationDetails)
    const selectedConversationRef = useRef<undefined | number>(undefined)
    const isGeneratingRef = useRef(false)
    const prevReconnectGenerationRef = useRef(0)
    const refetchTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null)
    /** Tracks the in-progress activity stage so we can roll it into history with a duration. */
    const activeStageRef = useRef<{ stage: string; label: string; startedAt: number } | null>(null)

    const clearStreamInterrupted = () => setStreamInterrupted(false)

    const clearActivity = () => {
        activeStageRef.current = null
        setActivityStage(null)
        setActivityLabel(null)
        setActivityDetail(null)
        setActivityHistory([])
    }

    async function GetConversationDetails(conversationId: number | undefined) {
        if (!conversationId) {
            setConversationDetails([])
            return
        }

        try {
            selectedConversationRef.current = conversationId
            const { data }: AxiosResponse<ConversationPromptType[]> =
                await axiosInstance.get(
                    `conversation_details?conversation_id=${conversationId}`,
                )

            const tempConversations: ConversationPrompt[] = data.map(
                (item) =>
                    new ConversationPrompt({
                        id: item.id,
                        message: item.message,
                        user_created: item.user_created,
                        created_timestamp: item.created_timestamp,
                        created: item.created,
                        tokens_in: item.tokens_in ?? null,
                        tokens_out: item.tokens_out ?? null,
                        citations: item.citations ?? [],
                        feedback: item.feedback ?? null,
                    }),
            )
            if (tempConversations.length === 1) {
                // first message still awaiting assistant reply placeholder
                tempConversations.push(
                    new ConversationPrompt({ message: '', user_created: false }),
                )
            }
            conversationRef.current = tempConversations
            setConversationDetails(tempConversations)
        } catch (err) {
            console.log('Failed to refetch conversation after reconnect', err)
        }
    }

    const schedulePostStreamRefetch = useCallback(() => {
        if (refetchTimerRef.current) clearTimeout(refetchTimerRef.current)
        const conversationId = selectedConversationRef.current
        if (!conversationId) return
        // Backend saves the assistant prompt after END (+ citations frame); brief delay
        // hydrates prompt ids, persisted citations, and feedback.
        refetchTimerRef.current = setTimeout(() => {
            void GetConversationDetails(conversationId)
        }, 800)
    }, [])

    useEffect(() => {
        GetConversationDetails(selectedConversation)
    }, [selectedConversation])

    useEffect(() => {
        return () => {
            if (refetchTimerRef.current) clearTimeout(refetchTimerRef.current)
        }
    }, [])

    // Streaming recovery: on reconnect after a drop mid-stream, refetch conversation
    // and mark the in-flight assistant turn interrupted (no resume protocol on backend).
    useEffect(() => {
        if (reconnectGeneration === prevReconnectGenerationRef.current) {
            return
        }
        prevReconnectGenerationRef.current = reconnectGeneration
        if (reconnectGeneration === 0) {
            return
        }

        const wasGenerating = isGeneratingRef.current || messageResponsePart.current !== 0
        if (!wasGenerating) {
            return
        }

        const partial = messageRef.current
        messageRef.current = ''
        messageResponsePart.current = 0
        isGeneratingRef.current = false
        setIsGeneratingMessage(false)
        setStateMessage('')
        setStreamInterrupted(true)
        clearActivity()

        const details = [...conversationRef.current]
        if (details.length > 0) {
            const last = details[details.length - 1]
            if (last && !last.user_created && (!last.message || last.message.length === 0)) {
                details[details.length - 1] = new ConversationPrompt({
                    message: partial
                        ? `${partial}\n\n${STREAM_INTERRUPTED_MESSAGE}`
                        : STREAM_INTERRUPTED_MESSAGE,
                    user_created: false,
                })
                conversationRef.current = details
                setConversationDetails(details)
            }
        }

        if (selectedConversationRef.current) {
            GetConversationDetails(selectedConversationRef.current)
        }
    }, [reconnectGeneration])

    useEffect(() => {
        /* register a consistent channel name for identifing this chat messages */
        const channelName = `ACCOUNT_ID_${account?.email}`
  
        /* subscribe to channel and register callback */        
        subscribe(channelName, (message: string) => {
            /* Versioned frames (citations #98, status #96) — ignore unknown types safely */
            const frame = parseVersionedFrame(message)
            if (frame) {
                if (frame.type === 'citations') {
                    const citations = parseCitationsData(frame.data)
                    const details = [...conversationRef.current]
                    for (let i = details.length - 1; i >= 0; i -= 1) {
                        if (!details[i].user_created && details[i].message) {
                            details[i] = new ConversationPrompt({
                                ...details[i],
                                citations,
                            })
                            break
                        }
                    }
                    conversationRef.current = details
                    setConversationDetails(details)
                    schedulePostStreamRefetch()
                } else if (frame.type === 'status') {
                    const statusData = parseStatusData(frame.data)
                    if (statusData) {
                        const now = Date.now()
                        const previous = activeStageRef.current
                        if (previous && previous.stage !== statusData.stage) {
                            const durationMs = now - previous.startedAt
                            setActivityHistory((history) => [
                                ...history,
                                {
                                    stage: previous.stage,
                                    label: previous.label,
                                    startedAt: previous.startedAt,
                                    finishedAt: now,
                                },
                            ])
                            trackEvent(AnalyticsEvents.ACTIVITY_STAGE_COMPLETED, {
                                stage: previous.stage,
                                durationMs,
                            })
                        }
                        activeStageRef.current = {
                            stage: statusData.stage,
                            label: statusData.label,
                            startedAt: previous && previous.stage === statusData.stage
                                ? previous.startedAt
                                : now,
                        }
                        setActivityStage(statusData.stage)
                        // Tokens take over the label once the model starts writing.
                        setActivityLabel(statusData.stage === 'writing' ? null : statusData.label)
                        setActivityDetail(statusData.detail ?? null)
                    }
                }
                // other unknown types: no-op
                return
            }

            if (message === 'END_OF_THE_STREAM_ENDER_GAME_42'){
                messageResponsePart.current = 0
            
                conversationRef.current.pop()
            
                const finalized = new ConversationPrompt({
                    message: `${messageRef.current}`,
                    user_created: false,
                })
                conversationRef.current = [...conversationRef.current, finalized]
                setConversationDetails([...conversationRef.current])
                messageRef.current = ''
                setStateMessage('')
                isGeneratingRef.current = false
                setIsGeneratingMessage(false)
                setStreamInterrupted(false)
                clearActivity()
                schedulePostStreamRefetch()
            }
            else if (message === 'START_OF_THE_STREAM_ENDER_GAME_42'){
                conversationRef.current = conversationDetails
                isGeneratingRef.current = true
                setIsGeneratingMessage(true)
                setStreamInterrupted(false)
                messageResponsePart.current = 2
  
            }else if (message === 'CONVERSATION_ID'){
                isGeneratingRef.current = true
                setIsGeneratingMessage(true)
                setStreamInterrupted(false)
                messageResponsePart.current = 1
            }else{
                if (messageResponsePart.current === 1){
                    isGeneratingRef.current = true
                    setIsGeneratingMessage(true)
                    // this has to do with the conversation id
                    if(!selectedConversation){
                        const conversationId = Number(message);
                        setSelectedConversation(conversationId)
                        selectedConversationRef.current = conversationId
                        trackEvent(AnalyticsEvents.CONVERSATION_CREATED, {
                          conversationId,
                        });
                    }   
                }
                else if (messageResponsePart.current === 2){
                    isGeneratingRef.current = true
                    setIsGeneratingMessage(true)
                    messageRef.current += message
                    setStateMessage(messageRef.current)
                }
                // ignore stray frames outside an active stream phase
            }
            
        })
  
        return () => {
            /* unsubscribe from channel during cleanup */
            unsubscribe(channelName)
        }
    }, [account, subscribe, unsubscribe, conversationDetails, selectedConversation, setSelectedConversation, schedulePostStreamRefetch])


    return(
        <MessageContext.Provider value={{
            stateMessage,
            setStateMessage,
            conversationDetails,
            setConversationDetails,
            isGeneratingMessage,
            streamInterrupted,
            clearStreamInterrupted,
            activityStage,
            activityLabel,
            activityDetail,
            activityHistory,
            clearActivity,
        }}>
            {children}
        </MessageContext.Provider>
    )

}

export { MessageContext, MessageProvider}
