import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Key, isKeyRelease, matchesKey } from "@earendil-works/pi-tui";
import { registerBashGuard, registerUserBashHardening } from "./bash/guard.ts";
import { registerWorkerControl } from "./control/inbox.ts";
import { acquireOrchestratorLock, describeHolder, releaseOrchestratorLock } from "./lock.ts";
import type { LockOwner } from "./proof-of-death.ts";
import { type AgiModeState, createModeController } from "./mode.ts";
import {
	AGI_SYSTEM_SUFFIX,
	CONTEXT_CUSTOM_TYPE,
	POST_COMPACTION_BLOCK,
	USER_TURN_SYSTEM_SUFFIX,
	contextPressureWarning,
	purgeContextBlocks,
} from "./prompt.ts";
import { StateStore, readIfExists, resolvePaths } from "./state.ts";
import { registerControlTool, registerSteerTool } from "./tools/control.ts";
import { registerNoteTool } from "./tools/note.ts";
import { registerSleepTool } from "./tools/sleep.ts";
import { registerWaitForAgentTool } from "./tools/wait.ts";
import { registerArchiveTool } from "./tools/state.ts";
import { registerWorkerTools } from "./tools/delegate.ts";
import { createFleetController } from "./worker/fleet.ts";
import { loadAllRuns } from "./worker/registry.ts";
import { Scheduler, WAKE_CUSTOM_TYPE, type WakePayload } from "./scheduler/index.ts";
import { loadWakes } from "./scheduler/wake.ts";
import { registerInspector } from "./ui/inspector.ts";
import { createSleepIndicator } from "./ui/sleep.ts";
import { registerWakeRenderer } from "./ui/renderers.ts";

const CONTEXT_WARN_PERCENT = 70;

interface SessionEntry {
	type: string;
	customType?: string;
	data?: unknown;
}

/** Last persisted agi-mode state on the current branch (R-CONF-3). */
function readPersistedState(ctx: ExtensionContext): AgiModeState | undefined {
	// getBranch(), not getEntries(): navigating away from a branch where AGI was
	// enabled must not restore that branch's mode into an unrelated one.
	const entries = ctx.sessionManager.getBranch() as SessionEntry[];
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i];
		if (entry === undefined) continue;
		if (entry.type === "custom" && entry.customType === "agi-mode") {
			return entry.data as AgiModeState | undefined;
		}
	}
	return undefined;
}

export default function agiExtension(pi: ExtensionAPI): void {
	// R-ROLE-1 (re-entrancy guard). This must remain the first statement of the
	// factory. A worker that registers orchestrator tools can delegate again,
	// producing recursive process fan-out. R-CONF-8: PI_AGI_DISABLE is an
	// unconditional kill switch for an extension that can spawn processes.
	const role = process.env.PI_AGI_ROLE;
	if (process.env.PI_AGI_DISABLE === "1" || role === "worker" || role === "off") {
		// Workers keep Pi's normal capabilities. Only the small control inbox is
		// installed; the orchestrator surface below remains structurally unreachable.
		if (process.env.PI_AGI_DISABLE !== "1") {
			if (role === "worker" && process.env.PI_AGI_RUN_DIR !== undefined) {
				registerWorkerControl(pi, process.env.PI_AGI_RUN_DIR);
			}
		}
		return;
	}

	const mode = createModeController(pi);
	registerBashGuard(pi, { isAgiEnabled: () => mode.isEnabled() });
	registerUserBashHardening(pi);
	let lockHeld = false;
	let lockDir: string | undefined;
	let observeOnly = false;
	let pendingCompactionNotice = false;
	let pendingUserTurnNotice = false;
	// F4: ctx must never be captured across an extension-instance replacement, but the
	// Scheduler's timers fire outside any handler and need one. It is refreshed by
	// every handler below and cleared on shutdown, so a stale ctx is never used.
	let liveCtx: ExtensionContext | undefined;
	let userHaltLatched = false;
	let interruptConfirmOpen = false;
	let terminalInputUnsubscribe: (() => void) | undefined;

	function sendWakeMessage(payload: WakePayload): boolean {
		try {
			pi.sendMessage(
				{
					customType: WAKE_CUSTOM_TYPE,
					// Extension-triggered turns do not emit `before_agent_start` in Pi.
					// The actionable event text is therefore the complete model input.
					content: payload.text,
					display: true,
					details: { reason: payload.reason },
				},
				{ triggerTurn: true, deliverAs: "followUp" },
			);
			return true;
		} catch {
			return false;
		}
	}

	const scheduler = new Scheduler({
		ctx: () => liveCtx,
		config: (ctx) => fleet.config(ctx),
		send: (payload) => {
			// R-UI-19 / F6: `custom_message`, because a wake payload must reach the model.
			// R-SLEEP-4: `followUp`, never `steer` — if the orchestrator happens to be
			// mid-turn we want the notification after it finishes its current reasoning,
			// not injected into the middle of it.
			return sendWakeMessage(payload);
		},
		onWakeReason: () => undefined,
		onRefresh: () => {
			if (liveCtx !== undefined) refreshSleepUI(liveCtx);
		},
		controlAvailable: () => fleet.supervisor() !== undefined,
	});

	const fleet = createFleetController({ onRunTerminal: (runId, status) => scheduler.onRunTerminal(runId, status) });
	const sleepIndicator = createSleepIndicator({ enabled: () => mode.isEnabled() && !observeOnly });

	function refreshSleepUI(ctx: ExtensionContext): void {
		if (!mode.isEnabled() || observeOnly) {
			sleepIndicator.restore(ctx);
			return;
		}
		if (scheduler.exitState() !== undefined) {
			sleepIndicator.restore(ctx);
			return;
		}
		if (!fleet.config(ctx).sleepIndicator) {
			sleepIndicator.restore(ctx);
			return;
		}
		sleepIndicator.setIdle(scheduler.isIdle());
		sleepIndicator.update(ctx, scheduler.indicatorState(ctx));
	}

	function store(ctx: ExtensionContext): StateStore {
		return new StateStore(ctx.cwd);
	}

	/**
	 * R-STATE-14. Refusal must happen at the door: two orchestrators sharing one
	 * runtime and worker fleet is the most destructive configuration
	 * available. On refusal the caller is offered read-only AGI.
	 */
	function claimLock(ctx: ExtensionContext): { ok: boolean; holder?: LockOwner; reason?: string } {
		if (lockHeld) return { ok: true };
		const paths = resolvePaths(ctx.cwd);
		const result = acquireOrchestratorLock(paths.orchestratorLock, ctx.sessionManager.getSessionId());
		if (!result.ok) return { ok: false, holder: result.holder, reason: result.reason };
		lockHeld = true;
		lockDir = paths.orchestratorLock;
		return { ok: true };
	}

	function releaseLock(): void {
		if (!lockHeld || lockDir === undefined) return;
		releaseOrchestratorLock(lockDir);
		lockHeld = false;
	}

	function enable(ctx: ExtensionContext, observe: boolean): void {
		liveCtx = ctx;
		userHaltLatched = false;
		if (observe) {
			observeOnly = true;
			mode.setEnabled(true, ctx);
			// The widget is read-only, so an observe session still gets it: seeing the
			// lock holder's fleet is the whole point of observe mode.
			fleet.activate(ctx);
			ctx.ui.notify("AGI mode enabled read-only (--observe): workers are visible, but delegation and archive lifecycle changes are refused.", "warning");
			return;
		}
		const claim = claimLock(ctx);
		if (!claim.ok) {
			// R-STATE-14 offers read-only supervision on refusal. Ordinary repository
			// file tools remain independent; worker and archive lifecycle operations do not.
			observeOnly = true;
			mode.setEnabled(true, ctx);
			fleet.activate(ctx);
			ctx.ui.notify(
				`AGI mode is read-only here: this working directory is already claimed by a live orchestrator (${describeHolder(claim.holder)}).\n` +
				`${claim.reason ?? ""}\nWorker delegation and archive lifecycle changes are refused.`,
				"warning",
			);
			return;
		}
		observeOnly = false;
		store(ctx).scaffold();
		mode.setEnabled(true, ctx);
		fleet.activate(ctx);
		// §10.7: adopt or orphan whatever is in the run directory before the
		// orchestrator's first turn, so startup recovery and completion wakes
		// reflect the current durable run state.
		fleet.reconcileOnStart(ctx);
		// R-SLEEP-5: re-arm whatever this session had armed before a restart, and
		// garbage-collect records belonging to other sessions (R-SLEEP-7).
		scheduler.activate(ctx);
		installEscHandler(ctx);
	}

	async function haltByUser(ctx: ExtensionContext, reason: string): Promise<boolean> {
		if (!mode.isEnabled()) {
			ctx.ui.notify("AGI mode is not enabled.", "warning");
			return false;
		}
		if (interruptConfirmOpen) return false;
		interruptConfirmOpen = true;
		let choice: string | undefined;
		try {
			choice = await ctx.ui.select("Interrupt the entire AGI run?", ["Yes — stop everything", "No — keep running"]);
		} finally {
			interruptConfirmOpen = false;
		}
		if (choice !== "Yes — stop everything") return false;

		// Latch and disarm before aborting. The abort emits settle/end events, and none
		// of those paths may re-arm a wake or resume the run after user intent is known.
		userHaltLatched = true;
		scheduler.deactivate();
		const supervisor = fleet.supervisor();
		const results = supervisor?.terminateOwned(reason, "user") ?? [];
		fleet.deactivate(ctx);
		releaseLock();
		observeOnly = false;
		mode.setEnabled(false, ctx);
		sleepIndicator.restore(ctx);
		terminalInputUnsubscribe?.();
		terminalInputUnsubscribe = undefined;
		ctx.abort();

		const stopped = results.filter((result) => result.state !== "skipped").length;
		ctx.ui.notify(`AGI run interrupted. ${stopped} owned worker${stopped === 1 ? "" : "s"} stopped; automatic resume is disabled.`, "info");
		return true;
	}

	function installEscHandler(ctx: ExtensionContext): void {
		terminalInputUnsubscribe?.();
		terminalInputUnsubscribe = undefined;
		if (!ctx.hasUI || typeof ctx.ui.onTerminalInput !== "function") return;
		terminalInputUnsubscribe = ctx.ui.onTerminalInput((data) => {
			if (!mode.isEnabled() || observeOnly || userHaltLatched || interruptConfirmOpen) return undefined;
			if (isKeyRelease(data) || !matchesKey(data, Key.escape)) return undefined;
			void haltByUser(ctx, "interrupted by the user with Esc");
			return { consume: true };
		});
	}

	async function disable(ctx: ExtensionContext): Promise<boolean> {
		// R-CTRL-18: only this orchestrator's own runs are in scope. `runs/` is shared
		// with every other pi session in the repo, and a terminal stop is unresumable,
		// so counting — or stopping — a concurrent session's work here would destroy it.
		const supervisor = fleet.supervisor();
		const active = supervisor === undefined
			? []
			: supervisor.partitionOwned(loadAllRuns(ctx.cwd).entries).owned;
		if (active.length > 0 && supervisor !== undefined) {
			const confirmed = await ctx.ui.confirm(
				"Turn AGI mode off and stop all workers?",
				`${active.length} active run(s) will be stopped terminally and cannot be resumed.`,
			);
			if (!confirmed) return false;
			// The confirmation promised a terminal stop, so the escalation that enforces it
			// runs here — before `fleet.deactivate` drops the pumps. Arming it through
			// stopAll and then detaching in the same call left a worker that never consumed
			// the durable request running indefinitely.
			supervisor.terminateOwned("AGI mode was turned off by the user", "user");
		}
		releaseLock();
		observeOnly = false;
		// R-SLEEP-3 / E64g: a wake outliving its purpose would wake an orchestrator
		// with nothing to do, so records are disarmed before the mode flag flips.
		scheduler.deactivate();
		fleet.deactivate(ctx);
		mode.setEnabled(false, ctx);
		// R-UI-9. Toggle OFF is the single most likely place to leak a custom working
		// indicator, and a leaked one breaks normal streaming for the rest of the
		// session. Restored last, after the mode flag is already false, so nothing can
		// re-install it behind this call.
		sleepIndicator.restore(ctx);
		terminalInputUnsubscribe?.();
		terminalInputUnsubscribe = undefined;
		return true;
	}

	pi.registerFlag("agi", {
		description: "Start in AGI orchestrator mode",
		type: "boolean",
		default: false,
	});

	registerWakeRenderer(pi);
	const noteEmitter = registerNoteTool(pi, {
		onLevel: (level) => {
			// R-ORCH-16 / R-TOOL-21: a blocked note suppresses both arming (§11.2) and
			// agi_sleep for the rest of the turn.
			if (level === "blocked") scheduler.noteBlocked();
		},
	});
	registerArchiveTool(pi, () => observeOnly);
	const controlDeps = { supervisor: () => fleet.supervisor(), readOnlySession: () => observeOnly };
	registerControlTool(pi, mode, controlDeps);
	registerSteerTool(pi, controlDeps);
	registerInspector(pi, { ...controlDeps, config: (ctx) => fleet.config(ctx), haltAll: haltByUser });
	registerSleepTool(pi, {
		config: (ctx) => fleet.config(ctx),
		enabled: () => mode.isEnabled() && !observeOnly,
		blocked: () => scheduler.isBlocked(),
		requestSleep: (request) => scheduler.requestSleep(request),
		activeRuns: (ctx) => scheduler.reviewEligibleRunCount(ctx),
		currentNote: () => scheduler.snapshot().sleepNote,
		noteEmitted: () => noteEmitter.hasEmitted(),
		emitNote: (text) => {
			noteEmitter.emit(text, "info");
		},
	});
	registerWaitForAgentTool(pi, {
		enabled: () => mode.isEnabled() && !observeOnly,
		blocked: () => scheduler.isBlocked(),
		activeRuns: (ctx) => scheduler.reviewEligibleRunCount(ctx),
		requestSleep: (request) => scheduler.requestSleep(request),
	});
	registerWorkerTools(pi, {
		supervisor: () => fleet.supervisor(),
		config: (ctx) => fleet.config(ctx),
		readOnlySession: () => observeOnly,
		onResultConsumed: (runId) => scheduler.onResultConsumed(runId),
		waitForWorker: () => scheduler.requestSleep({ source: "worker" }),
	});

	pi.registerCommand("agi", {
		description: "Toggle AGI orchestrator mode ('--observe' for read-only)",
		handler: async (args, ctx) => {
			const observe = args.trim() === "--observe";
		if (mode.isEnabled() && !observe) {
				await disable(ctx);
				return;
			}
			if (mode.isEnabled() && observe) {
				ctx.ui.notify("AGI mode is already on. Toggle it off first to re-enter read-only.", "warning");
				return;
			}
			enable(ctx, observe);
		},
	});

	pi.registerShortcut(Key.ctrlAlt("g"), {
		description: "Toggle AGI orchestrator mode",
		handler: async (ctx) => {
			if (mode.isEnabled()) {
				await disable(ctx);
				return;
			}
			enable(ctx, false);
		},
	});

	pi.registerCommand("agi-status", {
		description: "Show AGI mode, durable files, workers, lock, and scheduled wakes",
		handler: async (_args, ctx) => {
			const state = store(ctx);
			const describeFile = (label: string, file: string): string => {
				const status = state.fileStatus(file);
				if (!status.exists) return `${label}: missing`;
				return `${label}: ${status.bytes} bytes${status.error === undefined ? "" : ` (${status.error})`}`;
			};
			const runs = loadAllRuns(ctx.cwd).entries;
			const activeRuns = runs.filter((entry) => !["complete", "failed", "stopped", "timedOut", "orphaned"].includes(entry.status.state));
			const wakes = loadWakes(ctx.cwd, ctx.sessionManager.getSessionId()).own;
			const lines = [
				`mode: ${mode.isEnabled() ? (observeOnly ? "on (read-only)" : "on") : "off"}`,
				`state dir: ${state.paths.root}`,
				`orchestrator lock: ${lockHeld ? "held by this session" : "not held"}`,
				describeFile("goal.md", state.paths.goal),
				describeFile("plan.md", state.paths.plan),
				describeFile("memory/index.md", state.paths.memoryIndex),
				`notes/: ${state.listNotes().length} markdown file(s)`,
				`memory/: ${state.listMemoryFiles().length} markdown file(s), excluding index.md`,
				`workers: ${activeRuns.length} active, ${runs.length} total`,
				`scheduled wakes: ${wakes.length === 0 ? "none" : wakes.map((wake) => `${wake.kind}${wake.firesAt === undefined ? "" : ` at ${wake.firesAt}`}`).join(", ")}`,
			];
			ctx.ui.notify(lines.join("\n"), "info");
		},
	});

	pi.registerCommand("agi-goal", {
		description: "Edit .pi/agi/goal.md directly",
		handler: async (_args, ctx) => {
			const state = store(ctx);
			state.scaffold();
			// BUG-13: prefill the raw file, not a re-render of three parsed fields.
			// Re-rendering silently dropped every hand-added frontmatter field (an
			// `owner:` line, a comment) the moment the user saved — the opposite of
			// G8's promise that these files are hand-editable.
			const prefill = readIfExists(state.paths.goal) ?? "";
			const edited = await ctx.ui.editor("Edit AGI goal (goal.md)", prefill);
			if (edited === undefined || edited === prefill) return;
			try {
				const result = await state.writeGoal(edited);
				ctx.ui.notify(`goal.md saved (${result.bytesWritten} bytes).`, "info");
			} catch (error) {
				ctx.ui.notify(`goal.md not saved: ${(error as Error).message}`, "error");
			}
		},
	});

	pi.registerCommand("agi-memory", {
		description: "Browse and edit AGI memory files",
		handler: async (_args, ctx) => {
			const state = store(ctx);
			state.scaffold();
			const names = state.listMemories();
			if (names.length === 0) {
				ctx.ui.notify("No memory files yet. Create one with the normal write or edit tool.", "info");
				return;
			}
			const choice = await ctx.ui.select("AGI memory", names);
			if (choice === undefined) return;
			const memory = state.readMemory(choice);
			if (memory === undefined) return;
			const edited = await ctx.ui.editor(`Edit memory/${choice}.md`, memory);
			if (edited === undefined || edited === memory) return;
			try {
				const result = await state.writeMemory(choice, edited);
				ctx.ui.notify(`memory/${choice}.md saved (${result.bytesWritten} bytes).`, "info");
			} catch (error) {
				ctx.ui.notify(`memory/${choice}.md not saved: ${(error as Error).message}`, "error");
			}
		},
	});

	pi.on("tool_call", async (event, ctx) => {
		if (!mode.isEnabled()) return;
		liveCtx = ctx;
		// R-TOOL-21d / R-SLEEP-11: record action for tick backoff and for deciding at
		// settle whether the turn returned to the same external wait series.
		scheduler.noteAction(event.toolName);

	});

	// R-PROMPT-1/2: durable framing goes in the system prompt (cacheable); the
	// volatile state block is rebuilt from disk every turn so stale copies never
	// accumulate in history.
	pi.on("before_agent_start", async (event, ctx) => {
		liveCtx = ctx;
		if (!mode.isEnabled()) return;
		const userTurn = pendingUserTurnNotice;
		pendingUserTurnNotice = false;
		const state = store(ctx);
		const notices: string[] = [];
		// Scaffolding only creates directories and .gitignore. A failure is context,
		// never a gate on the model's normal tools or worker lifecycle.
		if (!observeOnly) {
			try { state.scaffold(); } catch (error) { notices.push(`State directories could not be prepared: ${(error as Error).message}`); }
		}
		if (!userTurn) {
			const recovery = fleet.takeRecoveryBlock();
			if (recovery !== undefined) notices.push(recovery);
		}
		if (pendingCompactionNotice) {
			notices.push(POST_COMPACTION_BLOCK);
			pendingCompactionNotice = false;
		}
		const usage = ctx.getContextUsage();
		const percent = usage?.percent ?? undefined;
		if (percent !== undefined && percent > CONTEXT_WARN_PERCENT) {
			notices.push(contextPressureWarning(percent));
		}
		if (observeOnly) {
			notices.push("This is a read-only supervision session. Use the lock-holding session for delegation and archive operations.");
		}
		const content = state.buildDigest(notices.length === 0 ? undefined : notices.join("\n\n")).content;

		return {
			systemPrompt: `${event.systemPrompt}${AGI_SYSTEM_SUFFIX}${userTurn ? USER_TURN_SYSTEM_SUFFIX : ""}`,
			message: { customType: CONTEXT_CUSTOM_TYPE, content, display: false },
		};
	});

	// R-UI-10: the extension only owns the working row between `agent_settled` and the
	// next `agent_start`. Taking it over while the orchestrator streams would replace
	// the real loader with a stale "supervising" line for the whole turn.
	pi.on("agent_start", async (_event, ctx) => {
		liveCtx = ctx;
		scheduler.onAgentStart();
		refreshSleepUI(ctx);
	});

	// R-SLEEP-2 / F2: `agent_settled`, never `agent_end`. `agent_end` fires during an
	// automatic retry and an auto-compaction retry too, so arming there fires a wake
	// into a turn pi was about to continue on its own.
	pi.on("agent_settled", async (_event, ctx) => {
		liveCtx = ctx;
		if (!mode.isEnabled() || observeOnly) return;
		try {
			scheduler.onAgentSettled(ctx);
		} catch (error) {
			// A scheduler fault must not propagate into pi's settle path, which would
			// leave the session mid-teardown. Surfaced, never fatal.
			ctx.ui.notify(`AGI scheduler: could not arm the next wake (${(error as Error).message}).`, "warning");
		}
		refreshSleepUI(ctx);
	});

	/**
	 * R-SLEEP-22 headless drain. Only for `!ctx.hasUI`, which is `print` and `json`:
	 * `hasUI` is true in RPC, so R-SLEEP-23 (RPC does not auto-drain, its client owns
	 * the lifecycle) is satisfied by the guard rather than by a separate branch.
	 */
	pi.on("agent_end", async (_event, ctx) => {
		liveCtx = ctx;
		if (!mode.isEnabled() || observeOnly || ctx.mode === "rpc") return;
		if (ctx.mode !== "print" && ctx.mode !== "json") return;
		const config = fleet.config(ctx);
		const supervisor = fleet.supervisor();
		const outcome = await scheduler.drainHeadless(ctx, {
			deadlineMs: config.headlessDrainMs,
			stopAll: (entries) => supervisor?.stopForHeadlessDrain(entries, "headless drain deadline reached"),
				deliver: (payload) => {
					if (!sendWakeMessage(payload)) throw new Error("headless wake delivery failed");
				},
		});
		if (outcome.timedOut) {
			// E68: never exit silently with orphans. This is an error, not a notice —
			// work was stopped mid-flight and the user has to know.
			ctx.ui.notify(
				`AGI headless drain hit its ${Math.round(config.headlessDrainMs / 60000)}m deadline with workers still running. They have been stopped. ` +
					`Inspect .pi/agi/.runtime/agents/ for partial results.`,
				"error",
			);
		}
	});

	/**
	 * R-SLEEP-19/20/21. The user's turn supersedes pending tick and sleep wakes; worker
	 * completion wakes are re-queued behind it rather than cancelled, because they
	 * report an independent event the orchestrator still has to learn about.
	 */
	pi.on("input", async (event, ctx) => {
		liveCtx = ctx;
		if (!mode.isEnabled() || observeOnly) return;
		if (event.text.length > 0) {
			try {
				await store(ctx).appendUserRequest(event.text);
			} catch {
				// The journal is a convenience reference. Turn handling stays independent.
			}
		}
		pendingUserTurnNotice = true;
		scheduler.onUserInput();
		// R-UI-9: the user is about to stream, so the working row goes back to pi.
		sleepIndicator.restore(ctx);
	});

	// R-PROMPT-2/3: keep only the newest block, and purge all of them when mode
	// is OFF so a disabled AGI mode leaves no residual instructions behind.
	pi.on("context", async (event) => {
		const messages = event.messages as AgentMessage[];
		const filtered = purgeContextBlocks(messages, mode.isEnabled());
		return filtered.length === messages.length ? undefined : { messages: filtered };
	});

	// R-COMPACT-2 cannot be implemented against pi 0.83. The spec says to append
	// to `customInstructions` on session_before_compact, but SessionBeforeCompactResult
	// has only {cancel, compaction}, and agent-session.ts passes its own local
	// `customInstructions` to compact() — a mutated event field is ignored. The
	// only supported override is returning a whole `compaction` result, which
	// R-COMPACT-2 explicitly rejects ("pi's summarizer is better"). Deferred
	// rather than faked; R-COMPACT-3 below still works and carries the same
	// authority statement into the post-compaction turn.
	pi.on("session_compact", async () => {
		pendingCompactionNotice = true;
	});

	pi.on("session_start", async (_event, ctx) => {
		liveCtx = ctx;
		mode.restore(readPersistedState(ctx), ctx);
		// F4: this is a fresh extension instance. Anything the previous one installed on
		// the working row belongs to a context that no longer exists (R-UI-9).
		sleepIndicator.restore(ctx);

		// R-CONF-7: flag values are only available at session_start, not in the factory.
		if (pi.getFlag("agi") === true && !mode.isEnabled()) {
			enable(ctx, false);
			return;
		}
		if (mode.isEnabled()) {
			const claim = claimLock(ctx);
			if (!claim.ok) {
				observeOnly = true;
				ctx.ui.notify(
					`AGI mode restored read-only: ${describeHolder(claim.holder)} holds the orchestrator lock. ${claim.reason ?? ""}`,
					"warning",
				);
				fleet.activate(ctx);
			} else {
				store(ctx).scaffold();
				fleet.activate(ctx);
				// §10.7 / R-CTRL-31: a restored AGI session is exactly the restart case.
				// Runs whose process is still alive are adopted, not killed — a detached
				// worker that survived the restart is still doing useful work.
				fleet.reconcileOnStart(ctx);
				// R-SLEEP-5 / E64f: a tick or sleep armed before the restart is re-armed,
				// and one whose firesAt has already passed fires immediately rather than
				// being silently extended.
				scheduler.activate(ctx);
				installEscHandler(ctx);
			}
		}
	});

	// R-CONF-3: tree navigation changes the active branch, so mode must be
	// recomputed from the new branch rather than left as the old branch's value.
	pi.on("session_tree", async (_event, ctx) => {
		liveCtx = ctx;
		mode.restore(readPersistedState(ctx), ctx);
		if (!mode.isEnabled()) {
			releaseLock();
			scheduler.deactivate();
			fleet.deactivate(ctx);
			// R-UI-9: navigating to a branch where AGI was never on must not leave a
			// supervising indicator behind.
			sleepIndicator.restore(ctx);
			terminalInputUnsubscribe?.();
			terminalInputUnsubscribe = undefined;
		} else if (!observeOnly) {
			installEscHandler(ctx);
		}
	});

	// R-UI-5: a badge surviving shutdown is a correctness bug, and this is the
	// path most often forgotten. R-CTRL-19: worker teardown here is best-effort and
	// bounded — a synchronous wait would risk hanging pi's exit, which is worse than
	// a lingering detached worker that the next session_start will reconcile.
	pi.on("session_shutdown", async (event, ctx) => {
		// R-SLEEP-5 / E70: the durable records deliberately survive a shutdown so the
		// next session_start re-arms them. `shutdown()` drops the in-process timers only;
		// `deactivate()` (which also deletes the records) is the toggle-OFF path.
		scheduler.shutdown();
		fleet.shutdown(ctx, event.reason);
		releaseLock();
		// R-UI-9: shutdown is one of the four named restore paths.
		sleepIndicator.restore(ctx);
		mode.clearUI(ctx);
		terminalInputUnsubscribe?.();
		terminalInputUnsubscribe = undefined;
		liveCtx = undefined;
	});
}