import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import type { StatusSnapshot } from "../domain/status.js";
import { resolvePendingImageMarkers, transformClipboardImages } from "../features/messages/image-input.js";
import {
	flushImagePreviewEntry,
	type ImagePreviewEntryData,
	registerImagePreviewSurface,
	stageImagePreviewData,
} from "../features/messages/image-preview.js";
import {
	beginAgentThoughtRun,
	finishAgentThoughtRun,
	rebuildAgentThoughtRunsFromEntries,
	refreshObservedThoughtComponents,
} from "../features/messages/thought-summary.js";
import { closeActiveBatch } from "../features/tools/boxed/batch.js";
import {
	beginAgentRun,
	finishAgentRun,
	invalidateTurnMembers,
	rebuildTurnRegistryFromEntries,
	registerTurnFromMessage,
	releaseTurnInvalidators,
} from "../features/tools/boxed/turn-summary.js";
import { requestToolPresentationRender } from "../features/tools/index.js";
import { registerPiStyleCommand } from "./commands.js";
import { type CompatibilityTestHooks, createPiStyleSessionCoordinator } from "./session-coordinator.js";
import { resetUsageFromSessionCache, usageFromSession } from "./session-usage.js";

/** Usage patch clears stale usage when the active branch/session has none. */
function usagePatch(ctx: ExtensionContext): StatusSnapshot {
	return { usage: usageFromSession(ctx.sessionManager) } as StatusSnapshot;
}

/**
 * Add Pi's read-only tools (grep/find/ls) to the active tool set if they are
 * registered. Preserves any other active tools (e.g. extension tools).
 * Only calls setActiveTools when something actually changed.
 */
function activateReadOnlyTools(pi: ExtensionAPI): void {
	const available = new Set(pi.getAllTools().map((tool) => tool.name));
	const active = new Set(pi.getActiveTools());
	let changed = false;
	for (const name of ["grep", "find", "ls"] as const) {
		if (available.has(name) && !active.has(name)) {
			active.add(name);
			changed = true;
		}
	}
	if (changed) pi.setActiveTools([...active]);
}

let compatibilityTestHooks: CompatibilityTestHooks = {};
export function __setCompatibilityTestHooks(hooks: CompatibilityTestHooks): () => void {
	const previous = compatibilityTestHooks;
	compatibilityTestHooks = hooks;
	return () => {
		compatibilityTestHooks = previous;
	};
}

/** Thin Pi adapter: register flags, commands, and forward lifecycle events. */
export default function piStyleExtension(pi: ExtensionAPI): void {
	// The core/message/tool surfaces are default-on (identity-certified per surface by
	// name/arity/source fingerprint, graceful native fallback for any surface whose
	// runtime identity is not recorded, conflict-preserving). The OFF switch is the
	// product gate `compatibility.allowCorePatches: false` (or `enabled: false`) in config.
	for (const [name, description] of [
		["pi-style-core-patches", "Enable pi-style message/tool core patches"],
		["pi-style-message-assistant", "Enable pi-style assistant message prefix"],
		["pi-style-message-special-blocks", "Enable pi-style boxed compaction/skill/branch/custom message blocks"],
		["pi-style-tools", "Enable pi-style tool renderer decoration"],
		["pi-style-readonly-tools", "Enable grep/find/ls read-only tools in the active tool set"],
	] as const)
		pi.registerFlag(name, { type: "boolean", description, default: true });
	// ASCII markers stay opt-in; unicode markers are the default.
	pi.registerFlag("pi-style-ascii", { type: "boolean", description: "Use ASCII pi-style markers" });
	const coordinator = createPiStyleSessionCoordinator(pi, compatibilityTestHooks);
	registerPiStyleCommand(pi, coordinator.app);
	// User-prompt image previews (ADR 0008): the entry renderer must be
	// registered before the first entry of this type can exist — the host drops
	// entries whose customType has no renderer — and once at extension load is
	// enough (persisted entries from previous sessions render on resume).
	registerImagePreviewSurface(pi);
	// User-prompt image previews (ADR 0008): the prompt's `images` are only
	// available at `before_agent_start`, but appending there would place the entry
	// above the user message because extension handlers run before the UI handles
	// `message_start(user)`. Capture there, then flush on the next event-loop turn
	// after `message_start(user)` so the preview appears immediately below it.
	// The assistant-start handler is a synchronous fallback for runtimes that begin
	// the assistant before the deferred flush runs. A steered prompt arriving before
	// flush replaces the staged slot (last write wins).
	let stagedImagePreview: ImagePreviewEntryData | undefined;
	let previewFlushTimer: ReturnType<typeof setTimeout> | undefined;
	pi.on("before_agent_start", (event) => {
		// Always replace the slot: a subsequent image-less prompt must not
		// accidentally flush the previous prompt's preview.
		stagedImagePreview = stageImagePreviewData(event.images ?? []);
	});
	pi.on("session_start", async (event, ctx) => {
		resetUsageFromSessionCache(ctx.sessionManager);
		// Pi only activates read/bash/edit/write by default; grep/find/ls are
		// registered but inactive (kept out of the model's tool list to keep the
		// core small). Activate them so the TUI shows them and the model can call
		// them directly, mirroring Claude Code's glob/grep/read tool set.
		if (pi.getFlag("pi-style-readonly-tools") === true) {
			activateReadOnlyTools(pi);
		}
		await coordinator.start(event, ctx);
	});
	pi.on("agent_start", () => {
		coordinator.app.runtime.current?.dismissStartup();
		// Thought segments are collected until agent_end (visible assistant text
		// splits them); tool summaries still use the whole agent-run boundary.
		beginAgentThoughtRun();
		beginAgentRun();
	});
	pi.on("input", async (event, _ctx) => {
		coordinator.app.runtime.current?.dismissStartup();
		// Bare `!`/`!!` submit guard: Pi treats `!`-prefixed input as a direct bash
		// command but falls through to normal message submission when the bang has
		// no command after it — sending a literal `!` to the agent. Drop those
		// accidental submits instead; Pi's submit path already cleared the editor
		// (onChange("") resets isBashMode), so the input returns to the normal
		// prompt without sending anything. Only the interactive input box is
		// guarded; rpc/extension sources keep sending text verbatim.
		if (event.source === "interactive") {
			const trimmed = event.text.trimStart();
			if (trimmed.startsWith("!")) {
				const bangLength = trimmed.startsWith("!!") ? 2 : 1;
				if (trimmed.slice(bangLength).trim() === "") return { action: "handled" };
			}
		}
		// Clipboard image input (ADR 0009): `[Image #N]` markers resolve for EVERY
		// submit source — a submit is a submit (image-paste semantics); consuming
		// the registry one-shot on rpc/extension submits too keeps a pasted marker
		// from dangling into a later, unrelated prompt. No-op when the registry is
		// empty. Raw clipboard path tokens, in contrast, only ever come from the
		// interactive editor's built-in paste — that upgrade stays interactive-only.
		// Combined images flow on to before_agent_start, where the ADR 0008 preview
		// entry picks them up.
		const markerResult = await resolvePendingImageMarkers(event.text);
		const pathResult = event.source === "interactive" ? await transformClipboardImages(event.text) : undefined;
		if (markerResult || pathResult) {
			return {
				action: "transform" as const,
				text: pathResult?.text ?? event.text,
				images: [...(markerResult?.images ?? []), ...(pathResult?.images ?? [])],
			};
		}
		return undefined;
	});
	pi.on("tool_execution_start", () => coordinator.app.runtime.current?.dismissStartup());
	pi.on("model_select", (event) =>
		coordinator.app.update(
			{
				model: event.model.name || event.model.id,
				...(event.model.provider ? { provider: event.model.provider } : {}),
				...(event.model.reasoning !== undefined ? { reasoning: event.model.reasoning } : {}),
			},
			"immediate",
		),
	);
	pi.on("thinking_level_select", (event) => coordinator.app.update({ thinkingLevel: event.level }, "immediate"));
	pi.on("session_info_changed", (event) =>
		coordinator.app.update({ sessionName: event.name }, "coalesced", { refreshExtensionStatuses: true }),
	);
	pi.on("message_start", (event) => {
		// A new message is a batch boundary: quiet-tool (read/ls/find) calls of the
		// new message start a fresh batch instead of joining the previous one.
		closeActiveBatch();
		if (event.message?.role === "user" && stagedImagePreview && !previewFlushTimer) {
			// Extension handlers run before Pi's interactive listener adds the user
			// message to the chat. Defer one turn so the entry is appended below the
			// visible user message, but do not wait for assistant content to arrive.
			const pending = stagedImagePreview;
			previewFlushTimer = setTimeout(() => {
				previewFlushTimer = undefined;
				if (stagedImagePreview !== pending) return;
				flushImagePreviewEntry(pi, pending);
				stagedImagePreview = undefined;
			}, 0);
		}
		// Fallback for runtimes that start the assistant before the deferred turn
		// runs. This preserves the ordering below the user and never duplicates.
		if (stagedImagePreview && event.message?.role === "assistant") {
			if (previewFlushTimer) {
				clearTimeout(previewFlushTimer);
				previewFlushTimer = undefined;
			}
			flushImagePreviewEntry(pi, stagedImagePreview);
			stagedImagePreview = undefined;
		}
		// grep/bash tree panels are NOT cleared here: historical panels must keep
		// their state so Pi re-renders of previous messages (scroll/resume) stay
		// intact. Only session boundaries reset them (session-coordinator).
	});
	pi.on("message_update", () => coordinator.app.update({}, "coalesced"));
	// Usage (tokens + cost) is aggregated from finalized session entries at
	// message/turn boundaries, mirroring Pi's native footer; per-chunk updates
	// stay usage-free to keep streaming cheap.
	pi.on("message_end", (_event, ctx) =>
		coordinator.app.update({ ...usagePatch(ctx) }, "coalesced", { refreshContextUsage: true }),
	);
	pi.on("turn_end", (event, ctx) => {
		// Append the finalized assistant message's tool batch to the current run.
		registerTurnFromMessage(event.message, event.toolResults);
		coordinator.app.update({ ...usagePatch(ctx) }, "deferred", { refreshContextUsage: true });
	});
	pi.on("agent_end", () => {
		// Finalize each contiguous thought segment before tool collapse. Rebuilding
		// bound assistant components leaves one clickable row per segment leader.
		finishAgentThoughtRun();
		// The run is complete: collapse its tool blocks into one summary line.
		// Pi only re-invokes the tool renderer selectors from updateDisplay(), so
		// the captured per-block invalidate callbacks force the collapse and the
		// captured Tui repaints. Interrupted runs (a call without a result) stay
		// expanded.
		const run = finishAgentRun();
		if (run) {
			invalidateTurnMembers(run);
			releaseTurnInvalidators(run);
			requestToolPresentationRender();
		}
	});
	pi.on("agent_settled", (_event, ctx) =>
		coordinator.app.update({ ...usagePatch(ctx) }, "coalesced", { refreshContextUsage: true }),
	);
	pi.on("session_tree", (_event, ctx) => {
		resetUsageFromSessionCache(ctx.sessionManager);
		// Rebuild both agent-run presentation registries from the selected branch.
		const entries = ctx.sessionManager.getEntries();
		rebuildAgentThoughtRunsFromEntries(entries);
		rebuildTurnRegistryFromEntries(entries);
		refreshObservedThoughtComponents();
		coordinator.app.update({ ...usagePatch(ctx) }, "deferred", { refreshContextUsage: true });
	});
	pi.on("session_compact", (_event, ctx) => {
		resetUsageFromSessionCache(ctx.sessionManager);
		coordinator.app.update({ ...usagePatch(ctx) }, "deferred", { refreshContextUsage: true });
	});
	pi.on("tool_result", (event, ctx) => {
		if (["write", "edit", "bash"].includes(event.toolName)) {
			coordinator.app.runtime.current?.invalidateGit();
			coordinator.app.update({ ...usagePatch(ctx) }, "delayed-retry", { refreshContextUsage: true });
		}
	});
	pi.on("user_bash", (_event, ctx) => {
		coordinator.app.runtime.current?.invalidateGit();
		coordinator.app.update({ ...usagePatch(ctx) }, "delayed-retry", { refreshContextUsage: true });
	});
	pi.on("session_shutdown", (_event, ctx) => {
		resetUsageFromSessionCache(ctx.sessionManager);
		coordinator.shutdown();
	});
}