import type { SessionManager } from "@earendil-works/pi-coding-agent";
import type { ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import { collectCompactionStats } from "./compaction-stats";
import { formatCheckpointCreatedMessage } from "./checkpoint-store";
import {
	buildCheckpointAssistantMessage,
	buildCheckpointDetails,
	buildRestoreConversationAssistantMessage,
	buildRestoreResultDetails,
	buildRestoreConversationToolResult,
	buildRestoreConfirmationMessage,
	createToolResultMessage,
} from "./messages";
import { CHECKPOINT_TOOL_NAME, RESTORE_CONVERSATION_TOOL_NAME } from "./constants";
import type { InvokedBy, PerformRollbackResult } from "./types";

export function getWritableSessionManager(ctx: ExtensionCommandContext): SessionManager {
	return ctx.sessionManager as unknown as SessionManager;
}

export function performCheckpoint(
	sm: SessionManager,
	checkpointId: number,
	toolCallId: string,
	label?: string,
): { toolResultEntryId: string; assistantEntryId: string } {
	const assistantEntryId = sm.appendMessage(buildCheckpointAssistantMessage(toolCallId, label));
	const details = buildCheckpointDetails(checkpointId, label);
	const firstUse = checkpointId === 1;
	const toolResultEntryId = sm.appendMessage(
		createToolResultMessage(
			toolCallId,
			CHECKPOINT_TOOL_NAME,
			formatCheckpointCreatedMessage(checkpointId, { firstUse }),
			details,
		),
	);

	return { toolResultEntryId, assistantEntryId };
}

export function performRollback(
	sm: SessionManager,
	checkpointEntryId: string,
	checkpointNumber: number,
	summary: string,
	invokedBy: InvokedBy,
	toolCallId: string,
	label?: string,
): PerformRollbackResult {
	const stats = collectCompactionStats(sm.getBranch(), checkpointEntryId);
	sm.branch(checkpointEntryId);

	const details = buildRestoreResultDetails(checkpointNumber, summary, invokedBy, stats, label);
	const assistantEntryId = sm.appendMessage(
		buildRestoreConversationAssistantMessage(toolCallId, checkpointNumber, summary, invokedBy),
	);
	const toolResultEntryId = sm.appendMessage(buildRestoreConversationToolResult(toolCallId, details));
	sm.appendMessage(buildRestoreConfirmationMessage(checkpointNumber, summary));

	return {
		assistantEntryId,
		toolResultEntryId,
		checkpointNumber,
	};
}

export async function forceNavigateToLeaf(ctx: ExtensionCommandContext, leafId: string): Promise<void> {
	const entry = ctx.sessionManager.getEntry(leafId);
	if (!entry) {
		throw new Error(`Entry ${leafId} not found`);
	}

	// Direct SessionManager writes advance the leaf without rebuilding agent state.
	// navigateTree() no-ops when already at the target, so pivot through the parent first.
	if (entry.parentId) {
		await ctx.navigateTree(entry.parentId);
	}
	await ctx.navigateTree(leafId);
}

/** @deprecated Use forceNavigateToLeaf */
export async function syncSessionLeaf(ctx: ExtensionCommandContext, leafId: string): Promise<void> {
	await forceNavigateToLeaf(ctx, leafId);
}