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

import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
	SessionEntry,
} from "@mariozechner/pi-coding-agent";
import { BorderedLoader, convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";
import { extractCandidateFiles, extractLoadedSkills, resolveExtractionModel } from "./session-analysis.js";
import { type HandoffExtraction, assembleHandoffDraft, extractHandoffContext } from "./handoff-extraction.js";
import { parseAndSetModelWithNotify } from "./model-utils.js";

export function registerHandoffCommand(pi: ExtensionAPI, startCountdown: (ctx: ExtensionContext) => void) {
	pi.registerCommand("handoff", {
		description: "Transfer context to a new session (-model provider/modelId)",
		handler: async (args: string, ctx: ExtensionCommandContext) => {
			if (!ctx.hasUI) {
				ctx.ui.notify("/handoff requires interactive mode", "error");
				return;
			}

			if (!ctx.model) {
				ctx.ui.notify("No model selected", "error");
				return;
			}

			// Parse optional -model flag from args
			let remaining = args;
			let newSessionModel: string | undefined;

			const modelMatch = remaining.match(/(?:^|\s)-model\s+(\S+)/);
			if (modelMatch) {
				newSessionModel = modelMatch[1];
				remaining = remaining.replace(modelMatch[0], " ");
			}

			let goal = remaining.trim();
			if (!goal) {
				const entered = await ctx.ui.input("handoff goal", "What should the new thread do?");
				if (!entered?.trim()) {
					ctx.ui.notify("Handoff cancelled", "info");
					return;
				}
				goal = entered.trim();
			}

			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) {
				ctx.ui.notify("No conversation to hand off", "warning");
				return;
			}

			const llmMessages = convertToLlm(messages);
			const conversationText = serializeConversation(llmMessages);
			const currentSessionFile = ctx.sessionManager.getSessionFile();
			const candidateFiles = [...extractCandidateFiles(branch, conversationText)];
			const loadedSkills = extractLoadedSkills(branch);
			const extractionModel = resolveExtractionModel(ctx) ?? ctx.model;

			const result = await ctx.ui.custom<HandoffExtraction | null>((tui, theme, _kb, done) => {
				const loader = new BorderedLoader(tui, theme, "Extracting handoff context...");
				loader.onAbort = () => done(null);

				const run = async () => {
					const auth = await ctx.modelRegistry.getApiKeyAndHeaders(extractionModel);
					if (!auth.ok) {
						throw new Error(`Failed to get API key: ${auth.error}`);
					}

					return extractHandoffContext(
						extractionModel,
						auth.apiKey,
						auth.headers,
						conversationText,
						goal,
						candidateFiles,
						loadedSkills,
						loader.signal,
					);
				};

				run()
					.then(done)
					.catch((err) => {
						console.error("handoff generation failed", err);
						done(null);
					});

				return loader;
			});

			if (!result) {
				ctx.ui.notify("Handoff cancelled", "info");
				return;
			}

			const prefillDraft = assembleHandoffDraft(result, goal, currentSessionFile);

			const editedPrompt = await ctx.ui.editor("Edit handoff draft", prefillDraft);
			if (editedPrompt === undefined) {
				ctx.ui.notify("Handoff cancelled", "info");
				return;
			}

			const next = await ctx.newSession({
				parentSession: currentSessionFile ?? undefined,
			});

			if (next.cancelled) {
				ctx.ui.notify("New session cancelled", "info");
				return;
			}

			// Apply -model if specified
			if (newSessionModel) {
				await parseAndSetModelWithNotify(newSessionModel, ctx, pi);
			}

			const newSessionFile = ctx.sessionManager.getSessionFile();
			ctx.ui.notify(`Switched to new session: ${newSessionFile ?? "(unnamed)"}`, "info");

			ctx.ui.setEditorText(editedPrompt);
			startCountdown(ctx);
		},
	});
}