import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { RUN_TIMER_ENTRY, type RunTimerEntry, RunClock } from "./clock.js";
import {
	ActiveSubagentWatch,
	querySubagentFleet,
	runIdFromPayload,
	SUBAGENT_ASYNC_COMPLETE,
	SUBAGENT_ASYNC_STARTED,
	SUBAGENT_PROCESS_TERMINAL,
} from "./subagents.js";

export { formatElapsed, RUN_TIMER_ENTRY, type RunTimerEntry, RunClock } from "./clock.js";

const WIDGET_KEY = "pi-run-timer";
const STATUS_KEY = "pi-run-timer";

type OptionalEventApi = {
	on(event: string, handler: (event: unknown, ctx: ExtensionContext) => unknown): void;
};

type ThemeLike = { fg(color: string, text: string): string };

type EntryApi = ExtensionAPI & {
	registerEntryRenderer?: (
		customType: string,
		renderer: (
			entry: { data?: RunTimerEntry },
			options: unknown,
			theme: ThemeLike,
		) => { render(): string[]; invalidate(): void },
	) => void;
};

type RenderTui = { requestRender(): void };

function onOptionalEvent(
	pi: ExtensionAPI,
	event: string,
	handler: (ctx: ExtensionContext) => void | Promise<void>,
) {
	(pi as OptionalEventApi).on(event, async (_event, ctx) => handler(ctx));
}

export default function (pi: ExtensionAPI) {
	const clock = new RunClock();
	let ctxRef: ExtensionContext | undefined;
	let tuiRef: RenderTui | undefined;
	let widgetBound = false;
	let recorded = false;
	const subagents = new ActiveSubagentWatch();
	const unsubscribes: Array<() => void> = [];

	(pi as EntryApi).registerEntryRenderer?.(RUN_TIMER_ENTRY, (entry, _options, theme) => {
		const label = entry.data?.label ?? "Worked • 0s";
		return {
			render: () => ["", theme.fg("muted", label)],
			invalidate() {},
		};
	});

	const paint = () => {
		if (!ctxRef?.hasUI) return;
		// Live chrome is only for an in-flight run. The settled time lives in
		// the transcript so it is not also drawn above the editor.
		if (!clock.isRunning) {
			ctxRef.ui.setWorkingMessage();
			ctxRef.ui.setStatus(STATUS_KEY, undefined);
			tuiRef?.requestRender();
			return;
		}
		const label = clock.display();
		ctxRef.ui.setWorkingMessage(label);
		ctxRef.ui.setStatus(STATUS_KEY, ctxRef.ui.theme.fg("muted", label));
		tuiRef?.requestRender();
	};

	const bindWidget = (ctx: ExtensionContext) => {
		if (!ctx.hasUI || widgetBound) return;
		widgetBound = true;
		ctx.ui.setWidget(WIDGET_KEY, (tui, theme) => {
			tuiRef = tui;
			return {
				render() {
					if (!clock.isRunning) return [];
					return [theme.fg("muted", clock.display())];
				},
				invalidate() {},
			};
		});
	};

	const restoreChrome = () => {
		if (!ctxRef?.hasUI) return;
		ctxRef.ui.setWorkingMessage();
		ctxRef.ui.setStatus(STATUS_KEY, undefined);
		ctxRef.ui.setWidget(WIDGET_KEY, undefined);
		widgetBound = false;
		tuiRef = undefined;
	};

	const record = () => {
		if (recorded) return;
		const snapshot = clock.snapshot();
		if (!snapshot) return;
		recorded = true;
		pi.appendEntry<RunTimerEntry>(RUN_TIMER_ENTRY, snapshot);
		if (!ctxRef?.hasUI) {
			console.error(snapshot.label);
		} else if (!(pi as EntryApi).registerEntryRenderer) {
			ctxRef.ui.notify(snapshot.label, "info");
		}
	};

	const keepLive = (ctx: ExtensionContext, at?: number) => {
		ctxRef = ctx;
		let created = false;
		if (clock.isRunning) {
			created = false;
		} else if (clock.isFrozen && (subagents.isActive || at !== undefined)) {
			created = clock.resume();
			if (created) recorded = false;
		} else {
			created = clock.start(at);
			if (created) recorded = false;
		}
		if (!ctx.hasUI) return;
		bindWidget(ctx);
		paint();
		if (created || clock.isRunning) clock.onTick(paint);
	};

	const finish = () => {
		if (subagents.isActive) {
			if (ctxRef) keepLive(ctxRef);
			return;
		}
		const justFroze = clock.freeze();
		if (!justFroze && recorded) return;
		record();
		paint();
	};

	const tryFinish = async () => {
		const fleet = await querySubagentFleet(pi.events);
		subagents.setRpcActive(fleet.totalActive);
		if (subagents.isActive) {
			if (ctxRef) keepLive(ctxRef, fleet.earliestStartedAt);
			return;
		}
		if (ctxRef?.isIdle() === false) return;
		finish();
	};

	const teardown = () => {
		for (const unsubscribe of unsubscribes.splice(0)) unsubscribe();
		if (clock.isRunning) clock.freeze();
		record();
		restoreChrome();
		clock.stop();
		recorded = false;
	};

	const listen = (event: string, handler: (payload: unknown) => void) => {
		const unsubscribe = pi.events.on(event, handler);
		if (typeof unsubscribe === "function") unsubscribes.push(unsubscribe);
	};

	listen(SUBAGENT_ASYNC_STARTED, (payload) => {
		subagents.noteStarted(runIdFromPayload(payload));
		if (ctxRef) keepLive(ctxRef);
	});
	listen(SUBAGENT_ASYNC_COMPLETE, (payload) => {
		subagents.noteComplete(runIdFromPayload(payload));
		void tryFinish();
	});
	listen(SUBAGENT_PROCESS_TERMINAL, (payload) => {
		subagents.noteComplete(runIdFromPayload(payload));
		void tryFinish();
	});

	pi.on("session_start", async (_event, ctx) => {
		ctxRef = ctx;
		const fleet = await querySubagentFleet(pi.events);
		subagents.setRpcActive(fleet.totalActive);
		if (!ctx.isIdle() || subagents.isActive) keepLive(ctx, fleet.earliestStartedAt);
	});

	pi.on("before_agent_start", async (_event, ctx) => {
		keepLive(ctx);
	});

	pi.on("agent_start", async (_event, ctx) => {
		keepLive(ctx);
	});

	pi.on("agent_end", async (_event, ctx) => {
		ctxRef = ctx;
		if (ctx.isIdle()) void tryFinish();
	});

	// Newer Pi emits agent_settled after retries/compaction. Older hosts ignore it.
	onOptionalEvent(pi, "agent_settled", (ctx) => {
		ctxRef = ctx;
		void tryFinish();
	});

	pi.on("session_shutdown", async (_event, ctx) => {
		ctxRef = ctx;
		teardown();
	});
}