import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

// pi-session-timer: per-run stats, exposed to the status row. One run =
// everything pi does to answer one of your messages. The formatted line is
// published to globalThis.__piSessionTimerLine so pi-alerts.ts can compose it
// onto the same status row as the alert/ponytail tags:
//   🔊 pi alert: idle 🔘 🦥 ponytail: ⚡ FULL ⌚ run 03:12 · 48.5K tok
// Session-level totals already live in the built-in status line; this only
// tracks the current/last run.
//
// Reset semantics: a user message_start arms the reset; it's applied when the
// next agent run actually starts (agent_start 0->1 transition). Auto-retries
// and auto-compact within an answer do NOT reset (no user message fires), and
// a message queued mid-run correctly scopes the NEXT run, not the current one.

function fmtClock(ms: number): string {
	const s = Math.floor(ms / 1000);
	if (s < 3600) return `${Math.floor(s / 60)}:${String(s % 60).padStart(2, "0")}`;
	return `${Math.floor(s / 3600)}:${String(Math.floor((s % 3600) / 60)).padStart(2, "0")}:${String(s % 60).padStart(2, "0")}`;
}

function fmtTok(n: number): string {
	if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
	if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`;
	return String(n);
}

export default function piSessionTimer(pi: ExtensionAPI) {
	let runAccumMs = 0;
	let runStart = 0;
	let activeRuns = 0; // depth counter: subagents/retries nest agent_start
	let tokIn = 0;
	let tokOut = 0;
	let sysPromptTokens: number | null = null; // footprint: first-message Usage.input (≈ system prompt + first user msg)
	let needReset = true; // first run of the session starts fresh
	let tick: ReturnType<typeof setInterval> | undefined;
	let uiCtx: any;

	function line(): string {
		const now = Date.now();
		const runMs = runAccumMs + (activeRuns > 0 ? now - runStart : 0);
		const label = activeRuns > 0 ? "run" : "last run";
		const sys = sysPromptTokens !== null ? ` · sys ~${fmtTok(sysPromptTokens)}` : "";
		return `⌚ ${label} ${fmtClock(runMs)} · ${fmtTok(tokIn + tokOut)} tok${sys}`;
	}

	function render() {
		const lineStr = line();
		(globalThis as any).__piSessionTimerLine = lineStr;
	}

	function startTick() {
		if (!tick) tick = setInterval(render, 1000);
	}

	function stopTick() {
		if (tick) {
			clearInterval(tick);
			tick = undefined;
		}
	}

	pi.on("session_start", async (_event, ctx) => {
		uiCtx = ctx;
		runAccumMs = 0;
		activeRuns = 0;
		tokIn = 0;
		tokOut = 0;
		sysPromptTokens = null;
		needReset = true;
		render();
	});

	pi.on("message_start", async (event) => {
		const m: any = event.message;
		if (m?.role === "user") needReset = true;
	});

	pi.on("agent_start", async (_event, ctx) => {
		uiCtx = ctx;
		if (activeRuns === 0) {
			if (needReset) {
				runAccumMs = 0;
				tokIn = 0;
				tokOut = 0;
				needReset = false;
			}
			runStart = Date.now();
		}
		activeRuns++;
		startTick();
		render();
	});

	pi.on("agent_end", async (_event, ctx) => {
		uiCtx = ctx;
		if (activeRuns > 0) {
			activeRuns--;
			if (activeRuns === 0) {
				runAccumMs += Date.now() - runStart;
				stopTick();
			}
		}
		render();
	});

	pi.on("message_end", async (event, ctx) => {
		uiCtx = ctx;
		const m: any = event.message;
		if (m?.role === "assistant" && m.usage) {
			// First assistant usage in the session ≈ system prompt + first user
			// message — the only cheap per-run signal of system-prompt footprint.
			if (sysPromptTokens === null) sysPromptTokens = m.usage.input ?? 0;
			tokIn += (m.usage.input ?? 0) + (m.usage.cacheRead ?? 0) + (m.usage.cacheWrite ?? 0);
			tokOut += m.usage.output ?? 0;
		}
		render();
	});

	pi.on("session_shutdown", async () => {
		stopTick();
	});
}