import { ref, nextTick, computed, watch } from 'vue';
import { useChatStore } from '../store/chat';
import { ws, agents } from '../store';
import { useSessionHistory } from './sessionHistory';
import { useMessages } from './useMessages';

export function useAgentSocket(
    visibleCount: { value: number },
    lastUsage: { value: any },
    pendingClearRef: { value: boolean },
    sentMessages: Set<string>,
    restoreLastSent?: () => void,
) {
    const chatStore = useChatStore();
    const { connected, send: wsSend, onMessage: onWsMessage, replayBuffer } = ws;
    const { updateFromServer, selectedAgent } = agents;

    // We strictly use chatStore actions for everything now
    const handoverInProgress = () => chatStore.smState === 'HANDOVER_PENDING' || chatStore.smState === 'HANDOVER_DONE';
    const isAgentRunning = () => chatStore.smState === 'AGENT_RUNNING';

    const history = useSessionHistory(isAgentRunning, visibleCount, () => selectedAgent.value);

    // Sync lastUsage back to caller
    history.lastUsage = lastUsage as any;

    // Keep chatStore.activeTurnCorrId in sync with the HUD-tracked active turn
    watch(history.activeTurnCorrId, (id) => { chatStore.activeTurnCorrId = id; });

    function pushSystem(text: string) {
        history.pushSystem(text);
    }

    function mount() {
        let wasReconnected = false;
        let wasJustSwitched = false;

        // ── Message handlers keyed by data.type ──────────────────────────
        // Each handler receives the raw WS message. Handlers close over
        // mount-scoped state (wasReconnected, wasJustSwitched) and module-
        // scoped refs (chatStore, history, etc.).

        const handlers: Record<string, (data: any) => void> = {
            auth_ok(data) { updateFromServer(data); },
            ready(data)   {
                updateFromServer(data);
                if (data.sessionId) chatStore.sessionKey = data.sessionId;
            },

            thinking(data) {
                if (!handoverInProgress()) chatStore.appendThinking(data.content);
            },

            delta(data) {
                if (!handoverInProgress()) {
                    history.flushPendingClear(pendingClearRef);
                    chatStore.collapseThinking();
                    chatStore.appendAssistantDelta(data.content, data.agentId);
                }
            },

            message(data) {
                if (handoverInProgress()) return;
                history.flushPendingClear(pendingClearRef);
                if (data.streaming === false) {
                    chatStore.createCompleteAssistantMessage(data.content, data.agentId, data.usage);
                } else if (data.final) {
                    chatStore.finalizeAssistantMessage(null, data.usage);
                }
                // streaming === true && !final handled by 'delta' to avoid double-text
            },

            truncated_warning(_data) {
                chatStore.collapseThinking();
                if (chatStore.hasActiveStreamingMessage()) chatStore.finalizeAssistantMessage(null, undefined, true);
                chatStore.truncatedWarning = true;
            },

            done(data) {
                if (handoverInProgress()) return;
                chatStore.collapseThinking();
                if (data.suppress) {
                    chatStore.suppressAssistantMessage();
                } else {
                    const doneContent: string | null = data.content || null;
                    if (chatStore.hasActiveStreamingMessage()) {
                        const deltaLen = chatStore.streamingMessageLength();
                        const useContent = (doneContent && deltaLen < doneContent.length) ? doneContent : null;
                        chatStore.finalizeAssistantMessage(useContent, data.usage);
                    } else if (doneContent) {
                        chatStore.createCompleteAssistantMessage(doneContent, undefined, data.usage);
                    }
                }
                history.lastSystemMsgRef.value = null;
            },

            session_history(data) {
                if (chatStore.hasActiveStreamingMessage()) chatStore.finalizeAssistantMessage(null);
                chatStore.collapseThinking();
                history.flushPendingClear(pendingClearRef);
                history.handleSessionHistory(data.entries);
            },

            hud(data) {
                history.pushHudEvent(data);
                if (data.event === 'turn_start' && !data.replay) {
                    chatStore.activeTurnCorrId = data.correlationId ?? null;
                    chatStore.startNewAssistantMessage(selectedAgent.value);
                }
            },

            event(data) {
                // Gateway agent stream events (tool, lifecycle, assistant)
                if (data.event === 'agent') {
                    const stream = data.payload?.stream;
                    if (stream === 'tool') {
                        console.log('[HUD agent/tool]', JSON.stringify(data.payload).slice(0, 400));
                    }
                }
            },

            tool(data) {
                // Legacy tool events — backward compat with older BE versions
                if (data.action === 'call') {
                    pushSystem(`${data.tool} ${data.args || ''}`);
                } else if (data.action === 'result') {
                    pushSystem(`→ ${data.result || ''}`);
                }
            },

            session_entry(data) {
                history.handleSessionEntry(data, sentMessages, pushSystem);
            },

            message_update(data) {
                // Patch an existing message by msgId (e.g. voice transcript, audio URL)
                if (!data.msgId || !data.patch) return;
                const patch = { ...data.patch };
                // Add auth token to voice audio URL
                if (patch.voiceAudioUrl) {
                    const token = localStorage.getItem('nyx_session') || localStorage.getItem('titan_token') || '';
                    patch.voiceAudioUrl = `${patch.voiceAudioUrl}?token=${encodeURIComponent(token)}`;
                }
                // Pre-add the gateway echo text to sentMessages so it gets deduped
                if (patch.transcript) {
                    sentMessages.add(`[voice transcript]: ${patch.transcript}`.trim());
                }
                const patched = chatStore.patchMessage(data.msgId, patch);
                if (!patched) {
                    console.warn('[message_update] no message found for msgId:', data.msgId);
                }
            },

            handover_done(data) {
                chatStore.messages.push({
                    role: 'assistant',
                    content: data.content || '📋 Handover written.',
                    agentId: selectedAgent.value,
                    sessionId: chatStore.localSessionId,
                });
                chatStore.messages.push({
                    role: 'system',
                    content: 'Start a new session with this handover context?',
                    confirmNew: true,
                    confirmed: false,
                    agentId: selectedAgent.value,
                    sessionId: chatStore.localSessionId,
                });
            },

            handover_context(_data) {
                // Silently discard — handover content already shown in chat history
            },

            // New two-SM protocol: channel_state + connection_state
            channel_state(data) {
                if (!data.state) return;
                if (data.clear_history) {
                    console.log('[clear] channel switch', { state: data.state, msgCount: chatStore.messages.length });
                    history.resetHudMaps();
                    chatStore.messages.length = 0;
                    pendingClearRef.value = false;
                }
                const prevChannel = chatStore.channelState;
                chatStore.applyChannelState(data.state);
                // On READY/FRESH: flush queued thought or clear after stop
                if (data.state === 'READY' || data.state === 'FRESH') {
                    history.lastSystemMsgRef.value = null;
                    if (chatStore.queuedThought !== null && prevChannel === 'AGENT_RUNNING') {
                        const thought = chatStore.queuedThought as string;
                        chatStore.queuedThought = null;
                        wsSend({ type: 'message', content: thought });
                    }
                }
            },

            connection_state(data) {
                if (!data.state) return;
                chatStore.applyConnectionState(data.state);
                if (data.state === 'LOADING_HISTORY') {
                    // Only stash on agent switch (clear_history comes via channel_state)
                    // Plain F5 reloads same session — no need to stash
                    history.resetHudMaps();
                    chatStore.messages.length = 0;
                    pendingClearRef.value = false;
                    wasReconnected = true;
                }
                if (data.state === 'SYNCED') {
                    history.flushPendingClear(pendingClearRef);
                    history.revealMessages();
                    wasReconnected = false;
                    wasJustSwitched = false;
                }
            },

            // Legacy: still handle session_state for backward compat
            session_state(data) {
                if (!data.state) return;
                if (data.reconnected) wasReconnected = true;
                if (data.reconnected || data.clear_history) {
                    console.log('[clear] immediate flush', { reconnected: data.reconnected, clear_history: data.clear_history, state: data.state, msgCount: chatStore.messages.length });
                    chatStore.stashMessages();
                    history.resetHudMaps();
                    chatStore.messages.length = 0;
                    pendingClearRef.value = false;
                }
                chatStore.applySessionState(data.state);
                if (data.state === 'READY' || data.state === 'FRESH' || data.state === 'IDLE') {
                    history.lastSystemMsgRef.value = null;
                    if (chatStore.queuedThought !== null) {
                        const thought = chatStore.queuedThought as string;
                        chatStore.queuedThought = null;
                        wsSend({ type: 'message', content: thought });
                    }
                }
            },

            session_total_tokens(data) {
                chatStore.sessionTotalTokens = data;
            },

            finance_update(data) {
                chatStore.finance = data;
            },

            usage(data) {
                if (!handoverInProgress()) {
                    chatStore.sessionTotalTokens = {
                        input_tokens: data.input_tokens || (chatStore.sessionTotalTokens?.input_tokens || 0),
                        cache_read_tokens: data.cache_read_tokens || (chatStore.sessionTotalTokens?.cache_read_tokens || 0),
                        output_tokens: data.output_tokens || (chatStore.sessionTotalTokens?.output_tokens || 0),
                    };
                }
            },

            session_status(data) {
                if (data.status === 'no_session') {
                    pendingClearRef.value = true;
                    history.flushPendingClear(pendingClearRef);
                    chatStore.messages.push({
                        role: 'system',
                        type: 'no_session',
                        content: '-- NO SESSION --',
                        agentId: selectedAgent.value,
                        sessionId: chatStore.localSessionId,
                    });
                    chatStore.sessionContextHint = '';
                    wasJustSwitched = false;
                    wasReconnected = false;
                } else if (data.status === 'watching') {
                    history.flushPendingClear(pendingClearRef);
                    history.sessionHistoryComplete.value = true;
                    history.revealMessages();
                    wasReconnected = false;
                    wasJustSwitched = false;
                }
            },

            sent(_data) {
                // User sent — transition handled by chatStore.send() in AgentsView
            },

            switch_ok(data) {
                wasJustSwitched = true;
                history.sessionHistoryComplete.value = false;
                if (data.sessionKey) chatStore.sessionKey = data.sessionKey;
            },

            new_ok(_data) {
                history.sessionHistoryComplete.value = false;
            },

            error(data) {
                if (data.code === 'SESSION_TERMINATED') {
                    chatStore.pushSystem('⚠️ Message not delivered — session was resetting. Please try again.', selectedAgent.value);
                    restoreLastSent?.();
                } else if (data.code === 'DISCARDED_NOT_IDLE' || data.code === 'DISCARDED_NOT_READY') {
                    chatStore.pushSystem('⚠️ Message not delivered — agent was busy. Please try again.', selectedAgent.value);
                    restoreLastSent?.();
                }
            },

            stopped(_data) {
                chatStore.pushSystem('✅ Agent stopped', selectedAgent.value);
            },

            killed(_data) {
                chatStore.pushSystem('☠️ Agent killed', selectedAgent.value);
            },
        };

        // ── Subscribe ────────────────────────────────────────────────────

        const unsubscribe = onWsMessage((data: any) => {
            const handler = handlers[data.type];
            if (handler) handler(data);
        });

        // Replay buffered messages through the same handler map
        replayBuffer((data: any) => {
            const handler = handlers[data.type];
            if (handler) handler(data);
        });

        return unsubscribe;
    }

    return {
        mount,
        lastSystemMsg: history.lastSystemMsgRef,
        hudTree: history.hudTree,
        hudVersion: history.hudVersion,
        getToolsForTurn: history.getToolsForTurn,
        hudSnapshot: history.hudSnapshot,
        toolCallMapSnapshot: history.toolCallMapSnapshot,
        sessionHistoryComplete: history.sessionHistoryComplete,
        pushSystem,
        hasActiveStreamingMessage: chatStore.hasActiveStreamingMessage,
    };
}
