import { completeSimple, type UserMessage } from "@earendil-works/pi-ai/compat";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Box, Markdown, Text, truncateToWidth, wrapTextWithAnsi, type TUI } from "@earendil-works/pi-tui";
import { getMarkdownTheme } from "@earendil-works/pi-coding-agent";

const ENTRY_TYPE = "btw-answer";
const WIDGET_ID = "btw-activity";
const SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
const SPINNER_INTERVAL_MS = 80;
const SYSTEM_PROMPT = `Answer the user's side question directly and concisely.
This is an isolated question: you have no access to the main Pi conversation and must not imply otherwise.
If essential context is missing, state the assumption or ask one short clarifying question.
Use Markdown when it improves readability.`;

interface BtwEntry {
	question: string;
	answer: string;
	model: string;
	timestamp: number;
}

export default function btwExtension(pi: ExtensionAPI) {
	let activeRequest: AbortController | undefined;
	let sessionActive = true;
	let spinnerTimer: ReturnType<typeof setInterval> | undefined;
	let spinnerIndex = 0;
	let activeTui: TUI | undefined;
	let activityAnswer = "";
	let activityDone = false;

	const stopSpinner = () => {
		if (spinnerTimer) clearInterval(spinnerTimer);
		spinnerTimer = undefined;
	};

	const stopActivity = (ctx?: ExtensionContext) => {
		stopSpinner();
		activeTui = undefined;
		activityAnswer = "";
		activityDone = false;
		ctx?.ui.setWidget(WIDGET_ID, undefined);
	};

	const finishActivity = (answer: string) => {
		stopSpinner();
		activityAnswer = answer;
		activityDone = true;
		activeTui?.requestRender();
	};

	const startActivity = (question: string, ctx: ExtensionContext) => {
		stopActivity(ctx);
		spinnerIndex = 0;
		activityAnswer = "";
		activityDone = false;
		const displayQuestion = question.replace(/\s+/g, " ");
		ctx.ui.setWidget(WIDGET_ID, (tui, theme) => {
			activeTui = tui;
			return {
				render(width: number) {
					const icon = activityDone
						? theme.fg("success", "✓")
						: theme.fg("accent", SPINNER_FRAMES[spinnerIndex] ?? SPINNER_FRAMES[0]);
					const command = theme.fg("accent", theme.bold("/btw"));
					const lines = [truncateToWidth(`${icon} ${command} ${theme.fg("muted", displayQuestion)}`, width)];
					if (activityAnswer) {
						const answerLines = wrapTextWithAnsi(activityAnswer, Math.max(1, width - 2));
						lines.push(...answerLines.slice(0, 10).map((line) => truncateToWidth(`  ${line}`, width)));
						if (answerLines.length > 10) lines.push(truncateToWidth(theme.fg("dim", "  …"), width));
					}
					return lines;
				},
				invalidate() {},
			};
		});
		spinnerTimer = setInterval(() => {
			spinnerIndex = (spinnerIndex + 1) % SPINNER_FRAMES.length;
			activeTui?.requestRender();
		}, SPINNER_INTERVAL_MS);
	};

	pi.registerEntryRenderer<BtwEntry>(ENTRY_TYPE, (entry, { expanded }, theme) => {
		const data = entry.data;
		if (!data) return new Text(theme.fg("error", "btw: invalid entry"), 0, 0);

		const box = new Box(1, 0, (text) => theme.bg("customMessageBg", text));
		box.addChild(
			new Text(
				`${theme.fg("accent", theme.bold("btw"))} ${theme.fg("muted", data.question)}`,
				0,
				0,
			),
		);
		box.addChild(new Markdown(data.answer, 0, 1, getMarkdownTheme()));
		if (expanded) {
			box.addChild(
				new Text(
					theme.fg("dim", `${data.model} · ${new Date(data.timestamp).toLocaleString()}`),
					0,
					0,
				),
			);
		}
		return box;
	});

	pi.registerCommand("btw", {
		description: "Ask an isolated side question without interrupting the main conversation",
		handler: async (rawArgs, ctx) => {
			const question = rawArgs.trim();

			if (question === "cancel") {
				if (!activeRequest) {
					ctx.ui.notify("No btw request is running", "info");
					return;
				}
				activeRequest.abort();
				return;
			}

			if (!question) {
				ctx.ui.notify("Usage: /btw <question>  ·  cancel with /btw cancel", "info");
				return;
			}

			if (activeRequest) {
				ctx.ui.notify("A btw request is already running · /btw cancel", "warning");
				return;
			}

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

			const model = ctx.model;
			const controller = new AbortController();
			activeRequest = controller;
			startActivity(question, ctx);

			// Deliberately detach from the command handler so Pi immediately returns
			// focus to the normal editor while the isolated completion keeps running.
			void (async () => {
				let keepCompletedActivity = false;
				try {
					const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
					if (!sessionActive || controller.signal.aborted) return;
					if (!auth.ok) throw new Error(auth.error);
					if (!auth.apiKey) throw new Error(`No API key for ${model.provider}`);

					const message: UserMessage = {
						role: "user",
						content: [{ type: "text", text: question }],
						timestamp: Date.now(),
					};
					const response = await completeSimple(
						model,
						{ systemPrompt: SYSTEM_PROMPT, messages: [message] },
						{
							apiKey: auth.apiKey,
							headers: auth.headers,
							env: auth.env,
							signal: controller.signal,
							cacheRetention: "none",
							reasoning: "minimal",
							maxTokens: 1200,
						},
					);

					if (!sessionActive) return;
					if (response.stopReason === "aborted" || controller.signal.aborted) {
						ctx.ui.notify("btw cancelled", "info");
						return;
					}
					if (response.stopReason === "error") {
						throw new Error(response.errorMessage || "Model request failed");
					}

					const answer = response.content
						.filter((part): part is { type: "text"; text: string } => part.type === "text")
						.map((part) => part.text)
						.join("\n")
						.trim();

					if (!answer) throw new Error(`Model stopped without an answer (${response.stopReason})`);

					pi.appendEntry<BtwEntry>(ENTRY_TYPE, {
						question,
						answer,
						model: `${model.provider}/${model.id}`,
						timestamp: Date.now(),
					});

					// Pi defers transcript rendering while its main inference is streaming.
					// Keep the answer visible in the widget until the transcript catches up.
					if (!ctx.isIdle()) {
						finishActivity(answer);
						keepCompletedActivity = true;
					}
				} catch (error) {
					if (!sessionActive) return;
					if (controller.signal.aborted) {
						ctx.ui.notify("btw cancelled", "info");
					} else {
						const message = error instanceof Error ? error.message : String(error);
						ctx.ui.notify(`btw failed: ${message}`, "error");
					}
				} finally {
					if (activeRequest === controller) activeRequest = undefined;
					if (sessionActive && !keepCompletedActivity) stopActivity(ctx);
				}
			})();
		},
	});

	pi.on("agent_settled", (_event, ctx) => {
		if (!activeRequest && activityDone) stopActivity(ctx);
	});

	pi.on("session_shutdown", (_event, ctx) => {
		sessionActive = false;
		activeRequest?.abort();
		activeRequest = undefined;
		stopActivity(ctx);
	});
}