import { ref } from 'vue';
import type { SubagentProcessEvent, SubagentToolExecution } from '../types/index.d';

type ToolExecutionBlockKind = 'args' | 'result' | 'error';

export function useChatExpansionState(options: {
	hasLongToolExecutionBlock: (execution: SubagentToolExecution, kind: ToolExecutionBlockKind) => boolean;
	hasProcessEventLongPayload: (event: SubagentProcessEvent) => boolean;
}) {
	const expandedStreamingToolResults = ref<Record<string, boolean>>({});
	const expandedSubagentEvidence = ref<Record<string, boolean>>({});
	const expandedSubagentProcesses = ref<Record<string, boolean>>({});
	const expandedSubagentProcessPayloads = ref<Record<string, boolean>>({});
	const expandedToolExecutionBlocks = ref<Record<string, boolean>>({});

	function toolExecutionBlockKey(execution: SubagentToolExecution, kind: ToolExecutionBlockKind) {
		return `${execution.key}:${kind}`;
	}

	function isToolExecutionBlockExpanded(execution: SubagentToolExecution, kind: ToolExecutionBlockKind) {
		return Boolean(expandedToolExecutionBlocks.value[toolExecutionBlockKey(execution, kind)]);
	}

	function toggleToolExecutionBlock(execution: SubagentToolExecution, kind: ToolExecutionBlockKind) {
		const key = toolExecutionBlockKey(execution, kind);
		expandedToolExecutionBlocks.value = {
			...expandedToolExecutionBlocks.value,
			[key]: !expandedToolExecutionBlocks.value[key],
		};
	}

	function isToolExecutionBlockCollapsed(execution: SubagentToolExecution, kind: ToolExecutionBlockKind) {
		return options.hasLongToolExecutionBlock(execution, kind) && !isToolExecutionBlockExpanded(execution, kind);
	}

	function isSubagentEvidenceExpanded(entryId: string) {
		return Boolean(expandedSubagentEvidence.value[entryId]);
	}

	function toggleSubagentEvidenceExpanded(entryId: string) {
		expandedSubagentEvidence.value = {
			...expandedSubagentEvidence.value,
			[entryId]: !expandedSubagentEvidence.value[entryId],
		};
	}

	function isSubagentProcessExpanded(entryId: string) {
		return Boolean(expandedSubagentProcesses.value[entryId]);
	}

	function toggleSubagentProcessExpanded(entryId: string) {
		expandedSubagentProcesses.value = {
			...expandedSubagentProcesses.value,
			[entryId]: !expandedSubagentProcesses.value[entryId],
		};
	}

	function processEventPayloadKey(entryId: string, index: number, event: SubagentProcessEvent) {
		return `${entryId}:${index}:${event.type}:${event.toolCallId || event.runId || ''}`;
	}

	function isProcessEventPayloadExpanded(entryId: string, index: number, event: SubagentProcessEvent) {
		return Boolean(expandedSubagentProcessPayloads.value[processEventPayloadKey(entryId, index, event)]);
	}

	function toggleProcessEventPayloadExpanded(entryId: string, index: number, event: SubagentProcessEvent) {
		const key = processEventPayloadKey(entryId, index, event);
		expandedSubagentProcessPayloads.value = {
			...expandedSubagentProcessPayloads.value,
			[key]: !expandedSubagentProcessPayloads.value[key],
		};
	}

	function isProcessEventPayloadCollapsed(entryId: string, index: number, event: SubagentProcessEvent) {
		return options.hasProcessEventLongPayload(event) && !isProcessEventPayloadExpanded(entryId, index, event);
	}

	function toggleStreamingToolExpanded(toolCallId: string) {
		expandedStreamingToolResults.value = {
			...expandedStreamingToolResults.value,
			[toolCallId]: !expandedStreamingToolResults.value[toolCallId],
		};
	}

	function isStreamingToolExpanded(toolCallId: string) {
		return Boolean(expandedStreamingToolResults.value[toolCallId]);
	}

	return {
		isProcessEventPayloadCollapsed,
		isProcessEventPayloadExpanded,
		isStreamingToolExpanded,
		isSubagentEvidenceExpanded,
		isSubagentProcessExpanded,
		isToolExecutionBlockCollapsed,
		isToolExecutionBlockExpanded,
		processEventPayloadKey,
		toggleProcessEventPayloadExpanded,
		toggleStreamingToolExpanded,
		toggleSubagentEvidenceExpanded,
		toggleSubagentProcessExpanded,
		toggleToolExecutionBlock,
		toolExecutionBlockKey,
	};
}
