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

import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import type { SessionManager } from "@mariozechner/pi-coding-agent";
import { Key, matchesKey, parseKey } from "@mariozechner/pi-tui";
import { registerHandoffCommand } from "./handoff-command.js";
import { registerHandoffTool, type PendingHandoff } from "./handoff-tool.js";
import { parseAndSetModelWithNotify } from "./model-utils.js";
import { registerSessionQueryTool } from "./session-query-tool.js";

const STATUS_KEY = "handoff";
const COUNTDOWN_SECONDS = 10;

type PendingAutoSubmit = {
	ctx: ExtensionContext;
	sessionFile: string | undefined;
	interval: ReturnType<typeof setInterval>;
	unsubscribeInput: () => void;
};

function statusLine(ctx: ExtensionContext, seconds: number): string {
	const accent = ctx.ui.theme.fg("accent", `handoff auto-submit in ${seconds}s`);
	const hint = ctx.ui.theme.fg("dim", "(type to edit, Esc to cancel)");
	return `${accent} ${hint}`;
}

export default function (pi: ExtensionAPI) {
	let pending: PendingAutoSubmit | null = null;
	let pendingHandoff: PendingHandoff | null = null;
	let handoffTimestamp: number | null = null;

	const clearPending = (ctx?: ExtensionContext, notify?: string) => {
		if (!pending) return;

		clearInterval(pending.interval);
		pending.unsubscribeInput();
		pending.ctx.ui.setStatus(STATUS_KEY, undefined);

		const local = pending;
		pending = null;

		if (notify && ctx) {
			ctx.ui.notify(notify, "info");
		} else if (notify) {
			local.ctx.ui.notify(notify, "info");
		}
	};

	const autoSubmitDraft = () => {
		if (!pending) return;

		const active = pending;
		const currentSession = active.ctx.sessionManager.getSessionFile();
		if (active.sessionFile && currentSession !== active.sessionFile) {
			clearPending(undefined);
			return;
		}

		const draft = active.ctx.ui.getEditorText().trim();
		clearPending(undefined);

		if (!draft) {
			active.ctx.ui.notify("Handoff draft is empty", "warning");
			return;
		}

		active.ctx.ui.setEditorText("");

		try {
			if (active.ctx.isIdle()) {
				pi.sendUserMessage(draft);
			} else {
				pi.sendUserMessage(draft, { deliverAs: "followUp" });
			}
		} catch {
			pi.sendUserMessage(draft);
		}
	};

	const startCountdown = (ctx: ExtensionContext) => {
		clearPending(ctx);

		let seconds = COUNTDOWN_SECONDS;
		ctx.ui.setStatus(STATUS_KEY, statusLine(ctx, seconds));

		const unsubscribeInput = ctx.ui.onTerminalInput((data) => {
			if (matchesKey(data, Key.escape)) {
				clearPending(ctx, "Handoff auto-submit cancelled");
				return { consume: true };
			}

			if (parseKey(data) !== undefined) {
				clearPending(ctx, "Handoff auto-submit stopped (editing)");
			}

			return undefined;
		});

		const interval = setInterval(() => {
			if (!pending) return;

			seconds -= 1;
			if (seconds <= 0) {
				autoSubmitDraft();
				return;
			}

			ctx.ui.setStatus(STATUS_KEY, statusLine(ctx, seconds));
		}, 1000);

		pending = {
			ctx,
			sessionFile: ctx.sessionManager.getSessionFile(),
			interval,
			unsubscribeInput,
		};
	};

	// --- Event handlers ---

	pi.on("session_before_switch", (_event, ctx) => {
		if (pending) clearPending(ctx);
	});

	pi.on("session_start", (event, ctx) => {
		if (pending) clearPending(ctx);
		// A proper session switch (e.g. /new) or fork fully resets agent state,
		// so clear the context filter to avoid hiding new messages.
		if (event.reason === "new" || event.reason === "resume" || event.reason === "fork") {
			handoffTimestamp = null;
		}
	});

	pi.on("session_before_fork", (_event, ctx) => {
		if (pending) clearPending(ctx);
	});

	pi.on("session_before_tree", (_event, ctx) => {
		if (pending) clearPending(ctx);
	});

	pi.on("session_tree", (_event, ctx) => {
		if (pending) clearPending(ctx);
	});

	pi.on("session_shutdown", (_event, ctx) => {
		if (pending) clearPending(ctx);
	});

	// --- Tool-path handoff coordination ---
	//
	// The /handoff command has ExtensionCommandContext with ctx.newSession()
	// which does a full agent reset. The handoff tool only gets
	// ExtensionContext, which lacks newSession(). So the tool stores a
	// pending handoff and these handlers complete it:
	//
	// 1. agent_end: after the agent loop finishes, switch sessions and
	//    send the handoff prompt in the next macrotask
	// 2. context: filter pre-handoff messages since the low-level session
	//    switch doesn't clear agent.state.messages
	// 3. session_start (above): clears the filter on proper switches

	pi.on("agent_end", (_event, ctx) => {
		if (!pendingHandoff) return;

		const { prompt, parentSession, newModel } = pendingHandoff;
		pendingHandoff = null;

		// Create new session first — this triggers session_start which clears
		// any stale handoffTimestamp. Then set the timestamp for the current
		// handoff so the context filter is active for the new session's first turn.
		(ctx.sessionManager as SessionManager).newSession({ parentSession });
		handoffTimestamp = Date.now();

		setTimeout(async () => {
			ctx.ui.setEditorText(prompt);
			startCountdown(ctx);

			if (newModel) {
				await parseAndSetModelWithNotify(newModel, ctx, pi);
			}
		}, 0);
	});

	pi.on("context", (event, _ctx) => {
		if (handoffTimestamp === null) return;

		const cutoff = handoffTimestamp;
		const newMessages = event.messages.filter((m) => m.timestamp >= cutoff);
		if (newMessages.length > 0) {
			return { messages: newMessages };
		}
	});

	// --- Register tools and commands ---

	registerSessionQueryTool(pi);
	registerHandoffTool(pi, (h) => {
		pendingHandoff = h;
	});
	registerHandoffCommand(pi, startCountdown);
}