import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { buildCheckpointToolResult } from "./messages";
import {
	formatCheckpointCreatedMessage,
	formatPickerLine,
	getNextCheckpointId,
	listCheckpointsOnBranch,
	resolveCheckpointOnBranch,
	validateSummary,
} from "./checkpoint-store";
import {
	cacheCommandActions,
	installCommandActionCapture,
	makeCommandContext,
} from "./command-actions";
import {
	CHECKPOINT_TOOL_NAME,
	RESTORE_CONVERSATION_TOOL_NAME,
	ROLLBACK_HANDLED_MESSAGE,
	ROLLBACK_MISSING_COMMAND_CONTEXT_MESSAGE,
} from "./constants";
import { renderRestoreCall, renderRestoreResult } from "./render";
import {
	getWritableSessionManager,
	performCheckpoint,
	performRollback,
	forceNavigateToLeaf,
} from "./session-ops";
import type { InvokedBy, RollbackArgs } from "./types";

installCommandActionCapture();

function withCachedCommandContext(
	ctx: ExtensionCommandContext,
	handler: (ctx: ExtensionCommandContext) => Promise<void>,
): Promise<void> {
	cacheCommandActions(ctx);
	return handler(ctx);
}

async function executeRollback(
	ctx: ExtensionCommandContext,
	params: {
		checkpointEntryId: string;
		checkpointNumber: number;
		checkpointLabel?: string;
		summary: string;
		invokedBy: InvokedBy;
		toolCallId: string;
	},
): Promise<void> {
	const sm = getWritableSessionManager(ctx);
	const result = performRollback(
		sm,
		params.checkpointEntryId,
		params.checkpointNumber,
		params.summary,
		params.invokedBy,
		params.toolCallId,
		params.checkpointLabel,
	);
	await forceNavigateToLeaf(ctx, result.toolResultEntryId);
}

export default function piCheckpointExtension(pi: ExtensionAPI) {
	pi.on("tool_call", async (event, ctx) => {
		if (event.toolName !== RESTORE_CONVERSATION_TOOL_NAME) {
			return;
		}

		const input = event.input as RollbackArgs;
		const summaryError = validateSummary(input.summary ?? "");
		if (summaryError) {
			return { block: true, reason: summaryError };
		}

		const resolved = resolveCheckpointOnBranch(ctx.sessionManager.getBranch(), input.checkpoint_id);
		if ("error" in resolved) {
			return { block: true, reason: resolved.error };
		}

		const commandCtx = makeCommandContext(ctx);
		if (!commandCtx) {
			return { block: true, reason: ROLLBACK_MISSING_COMMAND_CONTEXT_MESSAGE };
		}

		await executeRollback(commandCtx, {
			checkpointEntryId: resolved.entryId,
			checkpointNumber: resolved.info.id,
			checkpointLabel: resolved.info.label,
			summary: input.summary.trim(),
			invokedBy: "assistant",
			toolCallId: event.toolCallId,
		});

		return { block: true, reason: ROLLBACK_HANDLED_MESSAGE };
	});

	pi.registerTool({
		name: CHECKPOINT_TOOL_NAME,
		label: "Checkpoint",
		description:
			"Create a conversation-only context checkpoint. Filesystem is unchanged. Work after this point can later be compacted via restore_conversation(summary). Optional label for identification.",
		parameters: Type.Object({
			label: Type.Optional(Type.String({ description: "Optional label for this checkpoint" })),
		}),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const checkpointId = getNextCheckpointId(ctx.sessionManager.getEntries());
			return buildCheckpointToolResult(checkpointId, params.label);
		},
	});

	pi.registerTool({
		name: RESTORE_CONVERSATION_TOOL_NAME,
		label: "Restore Conversation",
		description:
			"Compact conversation context back to a checkpoint. Provide a summary of everything done since that checkpoint; that summary becomes your durable memory. Does not revert filesystem changes. Does not delete your findings — it replaces verbose intermediate messages with your summary.",
		executionMode: "sequential",
		parameters: Type.Object({
			checkpoint_id: Type.String({ description: "Checkpoint id, e.g. #3 or 3" }),
			summary: Type.String({
				description:
					"Summary of what happened since the checkpoint. Write in first person as your durable memory of that work.",
			}),
		}),
		async execute() {
			return {
				content: [{ type: "text" as const, text: ROLLBACK_HANDLED_MESSAGE }],
				details: { semantics: "conversation_compaction" as const },
			};
		},
		renderCall(args, theme) {
			return renderRestoreCall(
				{
					checkpoint_id: args.checkpoint_id,
					summary: args.summary,
					invokedBy: (args as RollbackArgs & { invokedBy?: InvokedBy }).invokedBy ?? "assistant",
				},
				theme,
			);
		},
		renderResult(result, _options, theme) {
			return renderRestoreResult(result, theme);
		},
	});

	pi.registerCommand("checkpoint", {
		description: "Create a context checkpoint (usage: /checkpoint [label])",
		handler: async (args, ctx) => {
			await withCachedCommandContext(ctx, async (commandCtx) => {
				await commandCtx.waitForIdle();
				const label = args.trim() || undefined;
				const sm = getWritableSessionManager(commandCtx);
				const checkpointId = getNextCheckpointId(sm.getEntries());
				const toolCallId = `call_user_checkpoint_${Date.now()}`;
				const result = performCheckpoint(sm, checkpointId, toolCallId, label);
				await forceNavigateToLeaf(commandCtx, result.toolResultEntryId);
				commandCtx.ui.notify(formatCheckpointCreatedMessage(checkpointId, { firstUse: checkpointId === 1 }), "info");
			});
		},
	});

	pi.registerCommand("restore_conversation", {
		description: "Restore conversation to a checkpoint",
		handler: async (_args, ctx) => {
			if (ctx.mode !== "tui" || !ctx.hasUI) {
				ctx.ui.notify("/restore_conversation requires interactive mode", "error");
				return;
			}

			await withCachedCommandContext(ctx, async (commandCtx) => {
				await commandCtx.waitForIdle();

				const checkpoints = listCheckpointsOnBranch(commandCtx.sessionManager.getBranch());
				if (checkpoints.length === 0) {
					commandCtx.ui.notify("No checkpoints on current branch", "warning");
					return;
				}

				const selected = await commandCtx.ui.select(
					"Restore conversation to checkpoint:",
					checkpoints.map((checkpoint) => formatPickerLine(checkpoint)),
				);
				if (!selected) {
					return;
				}

				const checkpoint = checkpoints.find((item) => formatPickerLine(item) === selected);
				if (!checkpoint) {
					return;
				}

				const summary = await commandCtx.ui.input("Summary to remember:", "What happened since this checkpoint...");
				const summaryError = validateSummary(summary ?? "");
				if (summaryError) {
					commandCtx.ui.notify(summaryError, "error");
					return;
				}

				await executeRollback(commandCtx, {
					checkpointEntryId: checkpoint.entryId,
					checkpointNumber: checkpoint.id,
					checkpointLabel: checkpoint.label,
					summary: summary!.trim(),
					invokedBy: "user",
					toolCallId: `call_user_restore_${Date.now()}`,
				});
			});
		},
	});
}