// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
// SPDX-FileCopyrightText: Petr Baudis <pasky@ucw.cz>
//
// SPDX-License-Identifier: MIT

import type { ExtensionAPI, SessionEntry } from "@mariozechner/pi-coding-agent";
import { convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";
import { extractCandidateFiles, extractLoadedSkills, resolveExtractionModel } from "./session-analysis.js";
import { type HandoffExtraction, assembleHandoffDraft, extractHandoffContext } from "./handoff-extraction.js";

export type PendingHandoff = {
	prompt: string;
	parentSession: string | undefined;
	newModel: string | undefined;
};

export function registerHandoffTool(pi: ExtensionAPI, setPendingHandoff: (h: PendingHandoff) => void) {
	pi.registerTool({
		name: "handoff",
		label: "Handoff",
		description:
			"Transfer context to a new session. Use when the user explicitly asks for a handoff or when the context window is nearly full. Provide a goal describing what the new session should focus on.",
		parameters: Type.Object({
			goal: Type.String({
				description: "The goal/task for the new session",
			}),
			model: Type.Optional(
				Type.String({
					description: "Model for the new session as provider/modelId (e.g. 'anthropic/claude-haiku-4-5')",
				}),
			),
		}),
		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
			if (!ctx.model) {
				return {
					content: [{ type: "text" as const, text: "No model selected." }],
					details: undefined,
				};
			}

			const branch = ctx.sessionManager.getBranch();
			const messages = branch
				.filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
				.map((entry) => entry.message);

			if (messages.length === 0) {
				return {
					content: [
						{
							type: "text" as const,
							text: "No conversation to hand off.",
						},
					],
					details: undefined,
				};
			}

			const llmMessages = convertToLlm(messages);
			const conversationText = serializeConversation(llmMessages);
			const candidateFiles = [...extractCandidateFiles(branch, conversationText)];
			const loadedSkills = extractLoadedSkills(branch);

			const extractionModel = resolveExtractionModel(ctx) ?? ctx.model;
			const auth = await ctx.modelRegistry.getApiKeyAndHeaders(extractionModel);
			if (!auth.ok) {
				return {
					content: [
						{
							type: "text" as const,
							text: `Failed to get API key: ${auth.error}`,
						},
					],
					details: undefined,
				};
			}

			let result: HandoffExtraction | null;
			try {
				result = await extractHandoffContext(
					extractionModel,
					auth.apiKey,
					auth.headers,
					conversationText,
					params.goal,
					candidateFiles,
					loadedSkills,
					signal,
				);
			} catch (err) {
				if (signal?.aborted) {
					return {
						content: [
							{
								type: "text" as const,
								text: "Handoff cancelled.",
							},
						],
						details: undefined,
					};
				}
				return {
					content: [
						{
							type: "text" as const,
							text: `Handoff extraction failed: ${String(err)}`,
						},
					],
					details: undefined,
				};
			}

			if (!result) {
				return {
					content: [
						{
							type: "text" as const,
							text: "Handoff extraction failed or was cancelled.",
						},
					],
					details: undefined,
				};
			}

			const currentSessionFile = ctx.sessionManager.getSessionFile();
			const prompt = assembleHandoffDraft(result, params.goal, currentSessionFile);

			setPendingHandoff({
				prompt,
				parentSession: currentSessionFile,
				newModel: params.model,
			});

			return {
				content: [
					{
						type: "text" as const,
						text: "Handoff prepared. The session will switch after this turn completes.",
					},
				],
				details: undefined,
			};
		},
	});
}