import { ref, triggerRef, nextTick } from 'vue';
import { useChatStore } from '../store/chat';

// ── HUD node types ────────────────────────────────────────────────────────────

export interface HudNode {
    id: string
    correlationId?: string
    type: 'turn' | 'tool' | 'think' | 'received'
    subtype?: string
    state: 'running' | 'done' | 'error'
    label: string
    tool?: string
    args?: Record<string, any>
    result?: Record<string, any>
    payload?: Record<string, any>
    startedAt: number
    endedAt?: number
    durationMs?: number
    children: HudNode[]
    replay: boolean
}

// ── Constants ─────────────────────────────────────────────────────────────────

const VISIBLE_PAGE = 50;
const MAX_HUD_NODES = 100;

// ── Composable ────────────────────────────────────────────────────────────────

export function useSessionHistory(
    isAgentRunning: () => boolean,
    visibleCount: { value: number },
    agentIdFn: () => string,
) {
    const store = useChatStore();
    const sessionHistoryComplete = ref(false);
    const lastUsage = ref<null | any>(null);
    let loadStartTime: number | null = null;
    let pendingMessages: any[] = [];
    let pendingUsageTotals: any | null = null;

    const lastSystemMsgRef = ref<string | null>(null);

    // ── HUD tree state ────────────────────────────────────────────────────────
    const hudTree = ref<HudNode[]>([]);
    const hudVersion = ref(0); // increments on every tree mutation — use as reactive dep in components
    // Map correlationId → running HudNode (for pairing _start/_end)
    const hudPending = new Map<string, HudNode>();
    // Map correlationId → HudNode (turns, for parenting tools into turns)
    const hudTurns = new Map<string, HudNode>();
    // Secondary index: toolCallId → WeakRef<HudNode> (OpenClaw-assigned ID, separate from corrId)
    // WeakRef: GC collects node when tree drops it; lazy eviction on access; clear on session reset.
    const toolCallMap = new Map<string, WeakRef<HudNode>>();

    function lookupByToolCallId(toolCallId?: string): HudNode | null {
        if (!toolCallId) return null;
        const ref = toolCallMap.get(toolCallId);
        if (!ref) return null;
        const node = ref.deref();
        if (!node) { toolCallMap.delete(toolCallId); return null; } // lazy eviction
        return node;
    }
    // Currently active (running) turn correlationId — stamped onto assistant messages
    const activeTurnCorrId = ref<string | null>(null);

    function makeNode(partial: Partial<HudNode>): HudNode {
        return {
            id: partial.id || crypto.randomUUID(),
            type: partial.type || 'received',
            state: partial.state || 'running',
            label: partial.label || '',
            children: [],
            replay: partial.replay ?? false,
            startedAt: partial.startedAt || Date.now(),
            ...partial,
        } as HudNode;
    }

    function findNode(nodes: HudNode[], corrId: string): HudNode | null {
        for (const n of nodes) {
            if (n.correlationId === corrId) return n;
            if (n.children) {
                const found = findNode(n.children, corrId);
                if (found) return found;
            }
        }
        return null;
    }

    function addHudNode(node: HudNode, parentCorrelationId?: string) {
        // 'history' is a sentinel parentId used by session-watcher for replay events with no real turn parent
        if (parentCorrelationId && parentCorrelationId !== 'history') {
            const parent = hudTurns.get(parentCorrelationId);
            if (parent) {
                parent.children.push(node);
                triggerRef(hudTree); hudVersion.value++;  // force Vue reactivity on nested mutation
                return;
            }
        }
        hudTree.value.unshift(node);
        if (hudTree.value.length > MAX_HUD_NODES) hudTree.value.splice(MAX_HUD_NODES);
        triggerRef(hudTree); hudVersion.value++;
    }

    function pushHudEvent(event: any) {
        const replay = !!event.replay;
        const ts: number = event.ts || Date.now();
        const corrId: string | undefined = event.correlationId;
        const parentId: string | undefined = event.parentId;

        switch (event.event) {
            case 'turn_start': {
                // Deduplicate: skip if a turn with this corrId already exists
                if (corrId && (hudPending.has(corrId) || hudTurns.has(corrId))) break;
                const node = makeNode({
                    type: 'turn', state: 'running',
                    label: '🔄 Turn',
                    correlationId: corrId,
                    startedAt: ts, replay,
                });
                if (corrId) { hudPending.set(corrId, node); hudTurns.set(corrId, node); }
                if (!replay && corrId) activeTurnCorrId.value = corrId;
                addHudNode(node);
                break;
            }
            case 'turn_end': {
                const node = corrId ? hudPending.get(corrId) : null;
                if (node) {
                    node.state = 'done';
                    node.endedAt = ts;
                    node.durationMs = event.durationMs;
                    if (corrId) { hudPending.delete(corrId); hudTurns.delete(corrId); }
                    triggerRef(hudTree); hudVersion.value++;
                }
                // else: no matching open turn — drop silently (no orphan node)
                if (!replay && corrId && activeTurnCorrId.value === corrId) activeTurnCorrId.value = null;
                break;
            }
            case 'think_start': {
                const node = makeNode({
                    type: 'think', state: 'running',
                    label: 'Thinking',
                    correlationId: corrId,
                    startedAt: ts, replay,
                });
                if (corrId) hudPending.set(corrId, node);
                addHudNode(node, parentId);
                break;
            }
            case 'think_end': {
                const node = corrId ? hudPending.get(corrId) : null;
                if (node) {
                    node.state = 'done';
                    node.endedAt = ts;
                    node.durationMs = event.durationMs;
                    if (corrId) hudPending.delete(corrId);
                    triggerRef(hudTree); hudVersion.value++;
                } else {
                    addHudNode(makeNode({ type: 'think', state: 'done', label: 'Thinking', correlationId: corrId, startedAt: ts, endedAt: ts, durationMs: event.durationMs, replay }), parentId);
                }
                break;
            }
            case 'tool_start': {
                const tool = event.tool || 'unknown';
                const args = event.args || {};
                const label = buildToolLabel(tool, args);
                const node = makeNode({
                    type: 'tool', state: 'running',
                    label, tool, args,
                    correlationId: corrId,
                    startedAt: ts, replay,
                });
                if (corrId) hudPending.set(corrId, node);
                // Register in toolCallMap for reliable tool_end pairing
                if (event.toolCallId) toolCallMap.set(event.toolCallId, new WeakRef(node));
                addHudNode(node, parentId);
                lastSystemMsgRef.value = label;
                break;
            }
            case 'tool_end': {
                const tool = event.tool || 'unknown';
                const result = event.result || {};
                // Lookup order: toolCallId (most reliable) → correlationId → tree scan fallback
                let node = lookupByToolCallId(event.toolCallId)
                        ?? (corrId ? hudPending.get(corrId) : null);
                // Last resort: find oldest running tool node under same turn
                if (!node && parentId) {
                    const turnNode = findNode(hudTree.value, parentId);
                    if (turnNode?.children) {
                        const match = turnNode.children.find(
                            n => n.type === 'tool' && n.state === 'running' && (!tool || n.tool === tool || n.tool === 'unknown')
                        );
                        if (match) {
                            node = match;
                            if (match.correlationId) hudPending.delete(match.correlationId);
                        }
                    }
                }
                if (node) {
                    node.state = result.ok === false ? 'error' : 'done';
                    node.result = result;
                    node.endedAt = ts;
                    node.durationMs = event.durationMs;
                    node.label = buildToolLabel(tool, node.args || {}, result);
                    if (corrId) hudPending.delete(corrId);
                    if (event.toolCallId) toolCallMap.delete(event.toolCallId);
                    triggerRef(hudTree); hudVersion.value++;
                } else {
                    // No matching start — create completed node (true orphan)
                    const label = buildToolLabel(tool, event.args || {}, result);
                    addHudNode(makeNode({ type: 'tool', state: 'done', label, tool, result, correlationId: corrId, startedAt: ts, endedAt: ts, durationMs: event.durationMs, replay }), parentId);
                }
                break;
            }
            case 'received': {
                const node = makeNode({
                    type: 'received', state: 'done',
                    subtype: event.subtype,
                    label: event.label || event.subtype || 'received',
                    startedAt: ts, endedAt: ts, replay,
                });
                addHudNode(node);
                break;
            }
        }
    }

    function buildToolLabel(tool: string, args: Record<string, any>, result?: Record<string, any>): string {
        const fileTools = ['read', 'write', 'edit', 'append'];
        if (fileTools.includes(tool)) {
            const vp: string = args.viewerPath || args.path || '';
            const filename = vp.split('/').pop() || vp;
            const area = result?.area || args.area;
            const areaStr = area ? `:L${area.startLine}–${area.endLine}` : '';
            return `${filename}${areaStr}`;
        }
        if (tool === 'exec') {
            const cmd: string = args.command || '';
            return `${cmd.slice(0, 60)}${cmd.length > 60 ? '…' : ''}`;
        }
        if (tool === 'web_fetch') return (args.url || '').slice(0, 60);
        if (tool === 'web_search') return (args.query || '').slice(0, 60);
        return tool;
    }

    // ── Turn → tools lookup ───────────────────────────────────────────────────

    function getToolsForTurn(corrId: string | null | undefined): HudNode[] {
        if (!corrId) return [];
        const turn = hudTree.value.find(n => n.correlationId === corrId)
            ?? [...hudTurns.values()].find(n => n.correlationId === corrId);
        return turn ? turn.children.filter(c => c.type === 'tool') : [];
    }

    // ── Debug snapshot ────────────────────────────────────────────────────────

    function hudSnapshot(): string {
        const lines: string[] = [`HUD tree — ${hudTree.value.length} root node(s)\n`];
        for (const node of hudTree.value) {
            const dur = node.durationMs != null ? ` [${node.durationMs}ms]` : '';
            const repl = node.replay ? ' (replay)' : '';
            lines.push(`  ${node.state === 'running' ? '⏳' : node.state === 'error' ? '❌' : '✅'} [${node.type}] ${node.label}${dur}${repl}`);
            lines.push(`     id=${node.id.slice(0,8)} corrId=${(node.correlationId || '—').slice(0,8)} children=${node.children.length}`);
            for (const child of node.children) {
                const cdur = child.durationMs != null ? ` [${child.durationMs}ms]` : '';
                lines.push(`       ${child.state === 'running' ? '⏳' : child.state === 'error' ? '❌' : '✅'} [${child.type}] ${child.label}${cdur}`);
                if (child.args) lines.push(`         args: ${JSON.stringify(child.args).slice(0, 80)}`);
                if (child.result) lines.push(`         result: ${JSON.stringify(child.result).slice(0, 80)}`);
            }
        }
        return lines.join('\n');
    }

    // ── Legacy pushSystem (kept for non-tool system messages) ────────────────

    function pushSystem(text: string) {
        lastSystemMsgRef.value = text;
    }

    // ── Pending clear ─────────────────────────────────────────────────────────

    function flushPendingClear(pendingClearRef: { value: boolean }) {
        if (!pendingClearRef.value) return;
        pendingClearRef.value = false;
        store.clearMessages();
        visibleCount.value = VISIBLE_PAGE;
        sessionHistoryComplete.value = false;
        loadStartTime = performance.now();
        pendingMessages = [];
        pendingUsageTotals = null;
    }

    // ── History reveal ────────────────────────────────────────────────────────

    function revealMessages() {
        loadStartTime = null;

        const _pending = pendingMessages;
        const _usage = pendingUsageTotals;
        pendingMessages = [];
        pendingUsageTotals = null;

        const idx = store.messages.findIndex(m => m.role === 'system' && m.content.includes('Loading session history...'));
        if (idx !== -1) {
            store.messages.splice(idx, 1, ..._pending);
        } else {
            store.messages.unshift(..._pending);
        }

        // Set session context hint based on current agent's messages only (_pending)
        const revealedCount = _pending.filter((m: any) => m.role !== 'system').length;
        store.sessionContextHint = revealedCount > 0 ? `${revealedCount} msgs in context` : 'fresh context';

        if (_usage) lastUsage.value = _usage;
    }

    // ── Bulk history handler ──────────────────────────────────────────────────

    function handleSessionHistory(entries: any[]) {
        if (!entries?.length) return;
        if (loadStartTime === null) loadStartTime = performance.now();

        if (!store.messages.some(m => m.content?.includes('Loading session history...'))) {
            store.pushSystem('⏳ Loading session history...', agentIdFn());
        }

        const newMsgs: any[] = [];
        const currentAgentId = agentIdFn();
        const currentSessionId = store.localSessionId;
        let pendingUsage: any = null;

        for (const data of entries) {
            // HUD events (hud protocol) — route to pushHudEvent
            if (data.type === 'hud') {
                pushHudEvent({ ...data, replay: true });
                continue;
            }
            if (data.event === 'tool_start' || data.event === 'tool_end' ||
                data.event === 'think_start' || data.event === 'think_end' ||
                data.event === 'turn_start' || data.event === 'turn_end' ||
                data.event === 'received') {
                pushHudEvent({ ...data, replay: true });
                continue;
            }

            if (data.entry_type === 'session_context') {
                newMsgs.push({ role: 'session_context', content: data.content || '', agentId: currentAgentId, sessionId: currentSessionId });
            } else if (data.entry_type === 'user_message') {
                newMsgs.push({ role: 'user', content: data.content || '', agentId: currentAgentId, sessionId: currentSessionId });
            } else if (data.entry_type === 'assistant_text') {
                const content = (data.content || '').replace(/^\[\[reply_to[^\]]*\]\]\s*/i, '').trim();
                if (!content) continue;
                const msg: any = { role: 'assistant', content, streaming: false, agentId: currentAgentId, sessionId: currentSessionId, timestamp: data.ts || null };
                if (data.truncated) msg.truncated = true;
                if (pendingUsage) { msg.usage = pendingUsage; pendingUsage = null; }
                newMsgs.push(msg);
            } else if (data.entry_type === 'usage') {
                pendingUsage = {
                    input_tokens:  data.input_tokens  || 0,
                    output_tokens: data.output_tokens || 0,
                    total_tokens:  data.total_tokens  || 0,
                    cost: Number(data.cost || 0),
                };
                const last = newMsgs[newMsgs.length - 1];
                if (last?.role === 'assistant') { last.usage = pendingUsage; pendingUsage = null; }
            }
        }
        pendingMessages = newMsgs;

        const totalUsage = entries
            .filter(e => e.entry_type === 'usage')
            .reduce((acc, e) => ({
                input_tokens:  acc.input_tokens  + (e.input_tokens  || 0),
                output_tokens: acc.output_tokens + (e.output_tokens || 0),
                total_tokens:  acc.total_tokens  + (e.total_tokens  || 0),
                cost: acc.cost + Number(e.cost || 0),
            }), { input_tokens: 0, output_tokens: 0, total_tokens: 0, cost: 0 });
        pendingUsageTotals = totalUsage.total_tokens > 0 ? totalUsage : null;
    }

    // ── Incremental entry handler (live + late-join) ──────────────────────────

    function handleSessionEntry(
        data: any,
        sentMessages: Set<string>,
        pushSystemFn: (text: string) => void,
    ) {
        // HUD events — route directly
        if (data.type === 'hud') { pushHudEvent(data); return; }

        const isReplay = !sessionHistoryComplete.value;
        const currentAgentId = agentIdFn();
        const currentSessionId = store.localSessionId;
        switch (data.entry_type) {
            case 'user_message': {
                const raw = data.content || '';
                if (raw.startsWith('A new session was started')) break;
                // Voice transcripts are already shown via message_update — skip echo
                if (!isReplay && raw.includes('[voice transcript]:')) break;
                // Dedup: check if we already have this message by msgId or content
                const hasByMsgId = data.msgId && store.messages.some((m: any) => m.msgId === data.msgId);
                if (hasByMsgId) break; // already shown via optimistic push
                if (!isReplay && !sentMessages.has(raw.trim())) {
                    store.messages.push({ role: 'user', content: raw, agentId: currentAgentId, sessionId: currentSessionId, msgId: data.msgId });
                } else {
                    sentMessages.delete(raw.trim());
                }
                break;
            }
            case 'assistant_text': break;
            case 'usage': break;
        }
    }

    function resetHudMaps() {
        hudPending.clear();
        hudTurns.clear();
        toolCallMap.clear();
        activeTurnCorrId.value = null;
        hudTree.value = [];
        hudVersion.value++;
    }

    function toolCallMapSnapshot(): Array<{ toolCallId: string; label: string | null; state: string | null; stale: boolean }> {
        return [...toolCallMap.entries()].map(([k, ref]) => {
            const node = ref.deref();
            return { toolCallId: k, label: node?.label ?? null, state: node?.state ?? null, stale: !node };
        });
    }

    return {
        sessionHistoryComplete,
        lastUsage,
        lastSystemMsgRef,
        hudTree,
        hudVersion,
        activeTurnCorrId,
        getToolsForTurn,
        pushHudEvent,
        hudSnapshot,
        toolCallMapSnapshot,
        resetHudMaps,
        flushPendingClear,
        revealMessages,
        handleSessionHistory,
        handleSessionEntry,
        pushSystem,
    };
}
