import type { AssistantMessage, ToolResultMessage } from "@earendil-works/pi-ai";
import { CHECKPOINT_TOOL_NAME, RESTORE_CONVERSATION_TOOL_NAME } from "./constants";
import { formatCheckpointCreatedMessage, formatCheckpointId, formatRestoreConversationResultMessage } from "./checkpoint-store";
import type { CheckpointDetails, CompactionStats, InvokedBy, RollbackArgs, RollbackResultDetails } from "./types";

const EMPTY_USAGE = {
	input: 0,
	output: 0,
	cacheRead: 0,
	cacheWrite: 0,
	totalTokens: 0,
	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};

export function createAssistantToolCallMessage(
	toolCallId: string,
	toolName: string,
	args: Record<string, unknown>,
): AssistantMessage {
	return {
		role: "assistant",
		content: [
			{
				type: "toolCall",
				id: toolCallId,
				name: toolName,
				arguments: args,
			},
		],
		api: "pi-checkpoint",
		provider: "pi-checkpoint",
		model: "pi-checkpoint",
		usage: EMPTY_USAGE,
		stopReason: "toolUse",
		timestamp: Date.now(),
	};
}

export function createToolResultMessage(
	toolCallId: string,
	toolName: string,
	text: string,
	details?: unknown,
	isError = false,
): ToolResultMessage {
	return {
		role: "toolResult",
		toolCallId,
		toolName,
		content: [{ type: "text", text }],
		details,
		isError,
		timestamp: Date.now(),
	};
}

export function buildCheckpointDetails(checkpointId: number, label?: string): CheckpointDetails {
	return {
		checkpointId,
		...(label ? { label } : {}),
		semantics: "conversation_bookmark",
		filesystemChanged: false,
	};
}

export function buildCheckpointToolResult(checkpointId: number, label?: string) {
	const details = buildCheckpointDetails(checkpointId, label);
	const firstUse = checkpointId === 1;

	return {
		content: [{ type: "text" as const, text: formatCheckpointCreatedMessage(checkpointId, { firstUse }) }],
		details,
	};
}

export function buildRestoreResultDetails(
	checkpointNumber: number,
	summary: string,
	invokedBy: InvokedBy,
	stats: CompactionStats,
	label?: string,
): RollbackResultDetails {
	return {
		checkpointId: checkpointNumber,
		...(label ? { label } : {}),
		invokedBy,
		compactedMessageCount: stats.compactedMessageCount,
		summaryAuthor: invokedBy,
		summary,
		semantics: "conversation_compaction",
		filesystemRestored: false,
		...(stats.compactedFrom ? { compactedFrom: stats.compactedFrom } : {}),
		...(stats.compactedTo ? { compactedTo: stats.compactedTo } : {}),
	};
}

export function buildRestoreConversationToolResult(
	toolCallId: string,
	details: RollbackResultDetails,
) {
	return createToolResultMessage(
		toolCallId,
		RESTORE_CONVERSATION_TOOL_NAME,
		formatRestoreConversationResultMessage(details),
		details,
		false,
	);
}

export function buildRestoreConfirmationMessage(
	checkpointNumber: number,
	summary: string,
): AssistantMessage {
	return {
		role: "assistant",
		content: [
			{
				type: "text",
				text: `I restored the conversation back to checkpoint #${checkpointNumber} with the following summary:\n\n${summary}`,
			},
		],
		api: "pi-checkpoint",
		provider: "pi-checkpoint",
		model: "pi-checkpoint",
		usage: EMPTY_USAGE,
		stopReason: "stop",
		timestamp: Date.now(),
	};
}

export function buildRestoreConversationToolCallArgs(
	checkpointNumber: number,
	summary: string,
	invokedBy: InvokedBy,
): RollbackArgs & { invokedBy: InvokedBy } {
	return {
		checkpoint_id: formatCheckpointId(checkpointNumber),
		summary,
		invokedBy,
	};
}

export function buildCheckpointAssistantMessage(toolCallId: string, label?: string) {
	return createAssistantToolCallMessage(
		toolCallId,
		CHECKPOINT_TOOL_NAME,
		label ? { label } : {},
	);
}

export function buildRestoreConversationAssistantMessage(
	toolCallId: string,
	checkpointNumber: number,
	summary: string,
	invokedBy: InvokedBy,
) {
	return createAssistantToolCallMessage(
		toolCallId,
		RESTORE_CONVERSATION_TOOL_NAME,
		{ ...buildRestoreConversationToolCallArgs(checkpointNumber, summary, invokedBy) },
	);
}