import {
	defineTool,
	type ExtensionAPI,
	type ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { closeDebug, dbg } from "./src/debug.ts";
import { type AfkWidget, createAfkWidget } from "./src/widget.ts";

const COMPLETION_TOOL = "afk_complete";
const WIDGET_KEY = "afk-creature";
const HUMAN_TOOLS = new Set(["ask_user_question", "mockup"]);
const MAX_EMPTY_CONTINUATIONS = 2;
const AWAY_PROMPT = `The user is away and expects you to get the current task and queued follow-ups done autonomously. Recover their intention from the full conversation and current worktree; normal queue order and latest instructions still apply. Choose reasonable interpretations instead of asking routine questions. Investigate failures, try materially different alternatives, and finish independent work rather than stopping at the first obstacle. Unavailable human interaction is not itself a blocker: use a sound non-interactive path. Do not invent answers, approvals, or extra tasks. When all requested work is finished, or only genuine unresolved blockers remain after useful alternatives are exhausted, call afk_complete, then give one concise final report of results, verification, consequential assumptions, and remaining blockers.`;

function setToolActive(pi: ExtensionAPI, active: boolean): void {
	const tools = pi.getActiveTools().filter((name) => name !== COMPLETION_TOOL);
	if (active) tools.push(COMPLETION_TOOL);
	pi.setActiveTools(tools);
}

export default function afkExtension(pi: ExtensionAPI): void {
	let active = false;
	let completed = false;
	let completionRequested = false;
	let emptyContinuations = 0;
	let sessionId: string | undefined;
	let pendingCompaction:
		| { sessionId: string; timer: ReturnType<typeof setTimeout> }
		| undefined;
	let widget: AfkWidget | undefined;
	let widgetUi: ExtensionContext["ui"] | undefined;

	const cancelCompaction = (): void => {
		if (!pendingCompaction) return;
		clearTimeout(pendingCompaction.timer);
		pendingCompaction = undefined;
	};
	pi.events.on("continuation:query", (data) => {
		if (!active || typeof data !== "object" || data === null) return;
		const query = data as { sessionId?: unknown; respond?: unknown };
		if (query.sessionId !== sessionId || typeof query.respond !== "function")
			return;
		// Keep this reply synchronous: emit does not await listeners.
		query.respond(completed ? "done" : "working");
	});

	const refreshWidget = (): void => widget?.invalidate();
	const clearWidget = (): void => {
		const oldWidget = widget;
		const oldUi = widgetUi;
		widget = undefined;
		widgetUi = undefined;
		oldWidget?.dispose();
		oldUi?.setWidget(WIDGET_KEY, undefined);
	};
	const showWidget = (ctx: ExtensionContext): void => {
		if (ctx.mode !== "tui") return;
		clearWidget();
		widgetUi = ctx.ui;
		ctx.ui.setWidget(
			WIDGET_KEY,
			(tui, theme) => {
				widget = createAfkWidget(tui, theme, () => completed);
				return widget;
			},
			{ placement: "belowEditor" },
		);
	};
	const disable = (): void => {
		cancelCompaction();
		sessionId = undefined;
		active = false;
		completed = false;
		completionRequested = false;
		emptyContinuations = 0;
		setToolActive(pi, false);
		clearWidget();
		dbg?.("mode", { outcome: "disabled" });
	};
	const enable = (ctx: ExtensionContext): void => {
		cancelCompaction();
		sessionId = ctx.sessionManager.getSessionId();
		active = true;
		completed = false;
		completionRequested = false;
		emptyContinuations = 0;
		setToolActive(pi, true);
		showWidget(ctx);
		dbg?.("mode", { outcome: "enabled" });
	};

	pi.registerTool(
		defineTool({
			name: COMPLETION_TOOL,
			label: "afk complete",
			description:
				"Finish AFK work after reporting results, evidence, assumptions, and genuine unresolved blockers.",
			parameters: Type.Object({}),
			defaultActive: false,
			exposure: "model-only",
			executionMode: "sequential",
			async execute() {
				if (active) completionRequested = true;
				return {
					content: [
						{
							type: "text",
							text: "AFK completion was requested. Give the user your final report.",
						},
					],
					details: undefined,
				};
			},
		}),
	);

	pi.registerCommand("afk", {
		description: "Keep working autonomously while I am away",
		async handler(_args, ctx) {
			if (active && !completed) return;
			enable(ctx);
			ctx.ui.notify(
				"AFK mode: Pi will continue until it completes the work.",
				"info",
			);
			pi.sendMessage(
				{
					customType: "afk-start",
					content: AWAY_PROMPT,
					display: false,
				},
				{ triggerTurn: true, deliverAs: "steer" },
			);
		},
	});
	pi.registerCommand("back", {
		description: "Return to normal Pi behavior",
		async handler(_args, ctx) {
			if (active) {
				pi.sendMessage(
					{
						customType: "afk-back",
						content:
							"The user is back. AFK delegation has ended; resume normal interaction and disregard previous AFK continuation instructions.",
						display: false,
					},
					{ triggerTurn: false },
				);
			}
			disable();
			ctx.ui.notify("Back to normal.", "info");
		},
	});

	pi.on("agent_start", cancelCompaction);
	pi.on("session_compact", (event, ctx) => {
		cancelCompaction();
		if (!active || completed || event.willRetry || ctx.hasPendingMessages())
			return;
		const candidate = {
			sessionId: ctx.sessionManager.getSessionId(),
			timer: setTimeout(() => {
				if (pendingCompaction !== candidate) return;
				pendingCompaction = undefined;
				if (
					!active ||
					completed ||
					sessionId !== candidate.sessionId ||
					ctx.sessionManager.getSessionId() !== candidate.sessionId ||
					!ctx.isIdle() ||
					ctx.hasPendingMessages()
				)
					return;
				try {
					pi.sendMessage(
						{
							customType: "afk-continue",
							content:
								"Compaction completed. Recover the current task from the compaction summary and worktree, then continue AFK work.",
							display: false,
						},
						{ triggerTurn: true, deliverAs: "followUp" },
					);
					dbg?.("continuation", { outcome: "continued" });
				} catch (error) {
					disable();
					dbg?.("continuation", { outcome: "error" });
					ctx.ui.notify(
						`AFK could not resume: ${error instanceof Error ? error.message : String(error)}`,
						"error",
					);
				}
			}, 0),
		};
		pendingCompaction = candidate;
	});
	pi.on("before_agent_start", (event) => {
		if (!active) return;
		event.systemPromptOptions.sections["afk"] = AWAY_PROMPT;
	});
	pi.on("message_start", (event) => {
		if (!active || event.message.role !== "user") return;
		completed = false;
		completionRequested = false;
		emptyContinuations = 0;
		refreshWidget();
	});
	pi.on("tool_execution_end", (event) => {
		if (active && !event.isError && !HUMAN_TOOLS.has(event.toolName))
			emptyContinuations = 0;
	});
	pi.on("tool_call", (event) => {
		if (!active || !HUMAN_TOOLS.has(event.toolName)) return;
		return {
			block: true,
			reason: `${event.toolName} is unavailable while the user is away. Continue without it or report a genuine blocker.`,
		};
	});
	pi.on("ui_prompt_start", (_event, ctx) => {
		if (!active) return;
		disable();
		ctx.abort();
		dbg?.("continuation", { outcome: "cancelled" });
	});
	pi.on("agent_before_settle", (event, ctx) => {
		if (!active) return;
		if (event.outcome !== "completed") {
			disable();
			return;
		}
		if (completionRequested) {
			completed = true;
			emptyContinuations = 0;
			refreshWidget();
			dbg?.("mode", { outcome: "completed" });
			return;
		}
		if (emptyContinuations >= MAX_EMPTY_CONTINUATIONS) {
			disable();
			ctx.ui.notify(
				"AFK stopped: model did not call afk_complete after repeated continuations.",
				"warning",
			);
			return {
				entries: [
					{
						type: "custom_message",
						customType: "afk-incomplete",
						content:
							"AFK stopped without afk_complete. The final report may be incomplete; review the last response and continue manually if work remains.",
						display: true,
					},
				],
			};
		}
		emptyContinuations += 1;
		dbg?.("continuation", { outcome: "continued" });
		return {
			entries: [
				{
					type: "custom_message",
					customType: "afk-continue",
					content:
						"Continue autonomous AFK work. Do not stop until you report the result and call afk_complete.",
					display: false,
				},
			],
			continue: true,
		};
	});
	pi.on("session_start", (_event, ctx) => {
		disable();
		dbg?.("session.start", { mode: ctx.mode });
	});
	pi.on("session_tree", () => disable());
	pi.on("session_shutdown", (_event, ctx) => {
		disable();
		dbg?.("session.shutdown", { mode: ctx.mode });
		closeDebug();
	});
}