import { defineStore } from 'pinia';
import { ref, computed, nextTick } from 'vue';

// Channel SM — shared by all users on the session
export type ChannelState = 'FRESH' | 'READY' | 'AGENT_RUNNING' | 'HANDOVER_PENDING' | 'HANDOVER_DONE' | 'RESETTING' | 'NO_SESSION';
// Connection SM — per-WS, per-user
export type ConnectionState = 'CONNECTING' | 'LOADING_HISTORY' | 'SYNCED' | 'SWITCHING';
// Legacy compat — includes STOP_PENDING (frontend-only)
export type SmState = ChannelState | ConnectionState | 'STOP_PENDING';

export interface FinanceData {
    nextTurnFloor: number;
    projectionDelta: number;
    currentContextTokens: number;
    lastTurnCost: number;
    pricing: { prompt: number, completion: number };
}

// Artifact types that belong to each pane
const DISPLAY_TYPES = new Set(['data_table', 'document_page', 'entity_detail']);
const KONSOLE_TYPES = new Set(['machine', 'action_bar', 'status']);

export const useChatStore = defineStore('chat', () => {
    // --- State ---
    const messages = ref<any[]>([]);

    // --- Artifacts ---
    const artifacts = ref<any[]>([]);
    const displayArtifacts = computed(() => artifacts.value.filter(a => DISPLAY_TYPES.has(a.type)));
    const konsoleArtifacts = computed(() => artifacts.value.filter(a => KONSOLE_TYPES.has(a.type)));

    function setArtifacts(arts: any[]) {
        artifacts.value = arts;
    }
    function clearArtifacts() {
        artifacts.value = [];
    }
    // Two-SM architecture: channel (shared) + connection (per-user)
    const channelState = ref<ChannelState>('NO_SESSION');
    const connectionState = ref<ConnectionState>('CONNECTING');
    // Legacy smState — derived from both, used by existing UI components
    const smState = computed<SmState>(() => {
        // Connection states take priority when not synced
        if (connectionState.value !== 'SYNCED') return connectionState.value;
        // When synced, show channel state (with STOP_PENDING override)
        return _stopPending.value ? 'STOP_PENDING' as SmState : channelState.value;
    });
    const _stopPending = ref(false);
    const truncatedWarning = ref<boolean>(false);
    const localSessionId = ref<string>(Math.random().toString(36).slice(2, 10));
    const sessionTotalTokens = ref<{ input_tokens: number; cache_read_tokens: number; output_tokens: number } | null>(null);
    const sessionUsage = ref<any>(null);
    const beVersion = ref<string>('');
    const finance = ref<FinanceData | null>(null);
    const sessionKey = ref<string>('');

    const sessionCost = computed<number>(() => {
        const p = finance.value?.pricing;
        const u = sessionTotalTokens.value;
        if (!p || !u) return 0;
        return ((u.input_tokens || 0) * p.prompt + (u.output_tokens || 0) * p.completion) / 1_000_000;
    });

    // Thought queue — 1 slot, replaced on each new queued send
    const queuedThought = ref<string | null>(null);

    // Active turn correlation id — set by sessionHistory, stamped onto assistant messages
    const activeTurnCorrId = ref<string | null>(null);

    // True when a handover confirm bubble is waiting for user action
    const handoverPending = computed<boolean>(() =>
        messages.value.some((m: any) => m.confirmNew && !m.confirmed)
    );

    const isRunning = computed<boolean>(() =>
        channelState.value === 'AGENT_RUNNING' ||
        _stopPending.value ||
        channelState.value === 'HANDOVER_PENDING'
    );

    // WS send fn — injected at app init by useAgentSocket/App.vue
    let _wsSend: ((payload: any) => void) | null = null;
    function setWsSend(fn: (payload: any) => void) { _wsSend = fn; }

    // Artifact action dispatch — sends {type:'action', action, data} over WS
    function sendAction(action: string, data: any = {}) {
        _wsSend?.({ type: 'action', action, data });
    }

    // Session actions (called from HudControls / HudRow)
    function newSession() {
        stashMessages();
        resetLocalSession();
        _wsSend?.({ type: 'new' });
    }

    function handover() {
        _wsSend?.({ type: 'handover_request' });
    }

    function stop() {
        if (channelState.value !== 'AGENT_RUNNING' && !_stopPending.value) return;
        queuedThought.value = null;
        _stopPending.value = true;
        pushSystem('Stopping after current turn...');
        _wsSend?.({ type: 'stop' });
    }

    function confirmNew() {
        const bubble = [...messages.value].reverse().find((m: any) => m.confirmNew && !m.confirmed);
        if (bubble) bubble.confirmed = true;
        stashMessages();
        resetLocalSession();
        _wsSend?.({ type: 'new' });
    }

    function stay() {
        const bubble = [...messages.value].reverse().find((m: any) => m.confirmNew && !m.confirmed);
        if (bubble) {
            bubble.confirmed = true;
            setTimeout(() => {
                const idx = messages.value.findIndex((m: any) => m === bubble);
                if (idx !== -1) messages.value.splice(idx, 1);
            }, 1000);
        }
        _wsSend?.({ type: 'cancel_handover' });
        // smState will be set by the backend's session_state:IDLE after cancel_handover
    }

    // --- SM write gatekeepers ---
    // Channel state — set by server broadcast to all users
    function applyChannelState(state: ChannelState) {
        channelState.value = state;
        if (state === 'AGENT_RUNNING') sessionContextHint.value = '';
        if (state === 'READY' || state === 'FRESH') _stopPending.value = false;
    }
    // Connection state — set by server for this user only
    function applyConnectionState(state: ConnectionState) {
        connectionState.value = state;
    }
    // Legacy: called by old code that sends a single state
    function applySessionState(state: SmState) {
        if (['FRESH', 'READY', 'AGENT_RUNNING', 'HANDOVER_PENDING', 'HANDOVER_DONE', 'NO_SESSION'].includes(state)) {
            applyChannelState(state as ChannelState);
        } else if (['CONNECTING', 'LOADING_HISTORY', 'SYNCED', 'SWITCHING'].includes(state)) {
            applyConnectionState(state as ConnectionState);
        }
    }
    // Called by App.vue on WS disconnect (pre-auth state)
    function setConnecting() {
        connectionState.value = 'CONNECTING';
    }

    // Streaming state
    let currentAssistantMessageIndex = -1;
    const streamingMessageVisibleContent = ref('');
    let streamingMessageFullContent = '';
    let streamingInterval: ReturnType<typeof setInterval> | null = null;
    let charIndex = 0;
    let pendingVisibleClear = false; // guards deferred nextTick clear

    // Thinking state
    let thinkingMessageIndex = -1;
    const thinkingContent = ref('');
    const sessionContextHint = ref('');
    const TYPING_TICK_MS = 10;

    // --- Getters ---
    const smLabel = computed<string>(() => {
        // Connection states take priority when not synced
        if (connectionState.value !== 'SYNCED') {
            switch (connectionState.value) {
                case 'CONNECTING':       return '⏳ Connecting...';
                case 'LOADING_HISTORY':  return '⏳ Loading...';
                case 'SWITCHING':        return '🔀 Switching...';
            }
        }
        if (_stopPending.value)          return '⛔ Stopping...';
        switch (channelState.value) {
            case 'AGENT_RUNNING':    return '⚙️ Working...';
            case 'HANDOVER_PENDING': return '📝 Handover...';
            case 'HANDOVER_DONE':    return '✅ Handover ready';
            case 'READY':            return sessionContextHint.value ? `✓ Ready - ${sessionContextHint.value}` : '● ✓ Ready';
            case 'FRESH':            return '✨ New session';
            case 'RESETTING':        return '🔄 Resetting...';
            case 'NO_SESSION':       return '○ No session';
            default:                 return '● ✓ Ready';
        }
    });

    const visibleMsgs = (visibleCount: number) => {
        return messages.value.slice(-visibleCount);
    };

    // --- Actions ---
    function resetLocalSession() {
        localSessionId.value = Math.random().toString(36).slice(2, 10);
        console.log('[ChatStore] Local Session Reset:', localSessionId.value);
        resetStreamingState();
    }

    /** Stash current messages to sessionStorage before clearing (survives F5) */
    function stashMessages() {
        const saveable = messages.value
            .filter((m: any) => m.role === 'user' || m.role === 'assistant')
            .map((m: any) => ({ role: m.role, content: m.content, agentId: m.agentId }));
        if (saveable.length > 0) {
            try { sessionStorage.setItem('hermes_prev_session', JSON.stringify(saveable)); }
            catch (_) { /* quota */ }
        }
    }

    function getPreviousSession(): any[] {
        try {
            const raw = sessionStorage.getItem('hermes_prev_session');
            return raw ? JSON.parse(raw) : [];
        } catch { return []; }
    }

    function clearMessages() {
        stashMessages();
        messages.value = [];
        sessionTotalTokens.value = null;
        finance.value = null;
        artifacts.value = [];
        resetStreamingState();
    }

    function pushMessage(msg: any) {
        const index = hasActiveStreamingMessage() ? currentAssistantMessageIndex : messages.value.length;
        const msgWithId = {
            ...msg,
            sessionId: localSessionId.value,
        };
        
        if (hasActiveStreamingMessage() && msg.role === 'user') {
            messages.value.splice(currentAssistantMessageIndex, 0, msgWithId);
            currentAssistantMessageIndex++;
        } else {
            messages.value.push(msgWithId);
        }
    }

    /** Patch an existing message by msgId — used for async updates (voice transcript, etc.) */
    function patchMessage(msgId: string, patch: Record<string, any>): boolean {
        const msg = messages.value.find((m: any) => m.msgId === msgId);
        if (!msg) return false;
        Object.assign(msg, patch);
        return true;
    }

    function pushSystem(text: string, agentId?: string) {
        const msg = {
            role: 'system',
            content: text,
            agentId: agentId || null,
            sessionId: localSessionId.value,
        };
        if (hasActiveStreamingMessage()) {
            messages.value.splice(currentAssistantMessageIndex, 0, msg);
            currentAssistantMessageIndex++;
        } else {
            messages.value.push(msg);
        }
    }

    // --- Streaming Actions ---
    function resetStreamingState() {
        if (streamingInterval !== null) clearInterval(streamingInterval);
        streamingInterval = null;
        currentAssistantMessageIndex = -1;
        streamingMessageFullContent = '';
        charIndex = 0;
        // Defer clearing visible content until after Vue renders msg.content,
        // preventing a blank frame. The flag lets a new stream cancel the clear.
        pendingVisibleClear = true;
        nextTick(() => { if (pendingVisibleClear) { streamingMessageVisibleContent.value = ''; pendingVisibleClear = false; } });
    }

    function startNewAssistantMessage(agentId?: string) {
        if (currentAssistantMessageIndex === -1 || !messages.value[currentAssistantMessageIndex]?.streaming) {
            pendingVisibleClear = false; // cancel any deferred clear from previous message
            resetStreamingState();
            messages.value.push({
                role: 'assistant',
                content: '',
                fullContent: '',
                usage: null,
                streaming: true,
                agentId: agentId ?? null,
                sessionId: localSessionId.value,
                turnCorrId: activeTurnCorrId.value ?? null,
            });
            currentAssistantMessageIndex = messages.value.length - 1;
        }
    }

    function appendAssistantDelta(delta: string, agentId?: string) {
        if (currentAssistantMessageIndex !== -1 && messages.value[currentAssistantMessageIndex]?.streaming) {
            streamingMessageFullContent += delta;
            messages.value[currentAssistantMessageIndex].fullContent = streamingMessageFullContent;
            if (!streamingInterval) startTypewriter();
        } else {
            startNewAssistantMessage(agentId);
            appendAssistantDelta(delta, agentId);
        }
    }

    function startTypewriter() {
        if (streamingInterval !== null) clearInterval(streamingInterval);
        streamingInterval = setInterval(() => {
            const backlog = streamingMessageFullContent.length - charIndex;
            if (backlog <= 0) {
                console.log('[hermes] typewriter done: charIndex=%d fullLen=%d visibleLen=%d streaming=%s',
                    charIndex, streamingMessageFullContent.length,
                    streamingMessageVisibleContent.value.length,
                    messages.value[currentAssistantMessageIndex]?.streaming);
                if (streamingInterval !== null) clearInterval(streamingInterval);
                streamingInterval = null;
                return;
            }
            const charsThisTick = backlog >= 200 ? 10 : backlog >= 50 ? 4 : 1;
            const end = Math.min(charIndex + charsThisTick, streamingMessageFullContent.length);
            streamingMessageVisibleContent.value += streamingMessageFullContent.slice(charIndex, end);
            charIndex = end;
        }, TYPING_TICK_MS);
    }

    function finalizeAssistantMessage(content?: string | null, usage?: any, truncated = false) {
        if (currentAssistantMessageIndex !== -1 && messages.value[currentAssistantMessageIndex]?.streaming) {
            const msg = messages.value[currentAssistantMessageIndex];
            const finalContent = ((content !== null && content !== undefined) ? content : streamingMessageFullContent)
                .replace(/\s*NO_REPLY\s*$/g, '').trim();
            console.log('[hermes] finalizeAssistantMessage: charIndex=%d fullLen=%d visibleLen=%d finalLen=%d',
                charIndex, streamingMessageFullContent.length,
                streamingMessageVisibleContent.value.length, finalContent.length);
            // Flush typewriter to end synchronously — prevents visible truncation
            // when done arrives while the typewriter interval is still mid-run
            if (streamingInterval !== null) {
                clearInterval(streamingInterval);
                streamingInterval = null;
            }
            streamingMessageVisibleContent.value = finalContent;
            charIndex = finalContent.length;
            msg.content = finalContent;
            msg.fullContent = finalContent;
            if (usage) msg.usage = usage;
            if (truncated) msg.truncated = true;
            msg.streaming = false;
            console.log('[hermes] post-finalize: msg.content.length=%d msg.streaming=%s idx=%d',
                msg.content.length, msg.streaming, currentAssistantMessageIndex);
            resetStreamingState();
        }
    }
    
    function appendThinking(content: string) {
        if (!content) return;
        thinkingContent.value += content;
        if (thinkingMessageIndex === -1) {
            messages.value.push({ role: 'thinking', content: thinkingContent, collapsed: false });
            thinkingMessageIndex = messages.value.length - 1;
        }
    }

    function collapseThinking() {
        if (thinkingMessageIndex !== -1 && messages.value[thinkingMessageIndex]) {
            messages.value[thinkingMessageIndex].collapsed = true;
        }
        thinkingMessageIndex = -1;
        thinkingContent.value = '';
    }

    function createCompleteAssistantMessage(content: string, agentId?: string, usage?: any) {
        // If there's an active streaming message, finalize it instead of creating a new one
        if (hasActiveStreamingMessage()) {
            finalizeAssistantMessage(content, usage);
            return;
        }
        
        const msg = {
            role: 'assistant',
            content: content.replace(/\s*NO_REPLY\s*$/g, '').trim(),
            fullContent: content.replace(/\s*NO_REPLY\s*$/g, '').trim(),
            usage: usage || null,
            streaming: false,
            agentId: agentId ?? null,
            sessionId: localSessionId.value,
            turnCorrId: activeTurnCorrId.value ?? null,
        };
        messages.value.push(msg);
    }

    function hasActiveStreamingMessage() {
        return currentAssistantMessageIndex !== -1 && messages.value[currentAssistantMessageIndex]?.streaming;
    }

    function streamingMessageLength() {
        return streamingMessageFullContent.length;
    }

    /**
     * Suppress (remove) the active streaming assistant message entirely.
     * Used for NO_REPLY turns: the agent emits NO_REPLY but partial deltas
     * already leaked to the browser. We drop the bubble retroactively on done.
     */
    function suppressAssistantMessage() {
        if (currentAssistantMessageIndex !== -1 && messages.value[currentAssistantMessageIndex]?.streaming) {
            if (streamingInterval !== null) {
                clearInterval(streamingInterval);
                streamingInterval = null;
            }
            messages.value.splice(currentAssistantMessageIndex, 1);
        }
        resetStreamingState();
    }

    return {
        messages,
        activeTurnCorrId,
        sessionKey,
        smState,   // legacy computed: connection priority, then channel
        channelState,
        connectionState,
        applySessionState,  // legacy compat
        applyChannelState,
        applyConnectionState,
        setConnecting,
        truncatedWarning,
        localSessionId, 
        sessionTotalTokens,
        sessionUsage,
        beVersion,
        sessionCost,
        finance,
        smLabel,
        sessionContextHint,
        streamingMessageVisibleContent,
        visibleMsgs,
        resetLocalSession,
        clearMessages,
        stashMessages,
        getPreviousSession,
        pushMessage,
        patchMessage,
        pushSystem,
        startNewAssistantMessage,
        appendAssistantDelta,
        finalizeAssistantMessage,
        createCompleteAssistantMessage,
        appendThinking,
        collapseThinking,
        hasActiveStreamingMessage,
        streamingMessageLength,
        suppressAssistantMessage,
        queuedThought,
        handoverPending,
        isRunning,
        setWsSend,
        sendAction,
        newSession,
        handover,
        stop,
        confirmNew,
        stay,
        artifacts,
        displayArtifacts,
        konsoleArtifacts,
        setArtifacts,
        clearArtifacts,
    };
});