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

type Totals = {
	input: number;
	output: number;
};

type AnimationState = {
	ctx?: ExtensionContext;
	visible: boolean;
	committed: Totals;
	target: Totals;
	from: Totals;
	displayed: Totals;
	startedAt: number;
	liveInputEstimate: number;
	liveOutputEstimate: number;
	outputCharRemainder: number;
	timer?: ReturnType<typeof setInterval>;
};

const WIDGET_ID = "token-odometer";
const FRAME_MS = 50;
const DURATION_MS = 900;

// These are intentionally conservative: the final provider usage usually lands
// above the live estimate, so the odometer mostly rolls upward and then settles
// on the exact count when usage arrives.
const INPUT_CHARS_PER_TOKEN = 6;
const OUTPUT_CHARS_PER_TOKEN = 6;

function easeOutCubic(t: number): number {
	const clamped = Math.max(0, Math.min(1, t));
	return 1 - Math.pow(1 - clamped, 3);
}

function getUsageNumber(value: unknown): number {
	return typeof value === "number" && Number.isFinite(value) ? Math.max(0, Math.floor(value)) : 0;
}

function add(a: Totals, b: Totals): Totals {
	return { input: a.input + b.input, output: a.output + b.output };
}

function usageFromMessage(message: any): Totals {
	const usage = message?.usage;
	return {
		input: getUsageNumber(usage?.input),
		output: getUsageNumber(usage?.output),
	};
}

function readSessionTotals(ctx: ExtensionContext): Totals {
	const totals: Totals = { input: 0, output: 0 };

	for (const entry of ctx.sessionManager.getBranch()) {
		if (entry.type !== "message") continue;
		const message = entry.message;
		if (message?.role !== "assistant") continue;

		const usage = message.usage;
		if (!usage) continue;

		totals.input += getUsageNumber(usage.input);
		totals.output += getUsageNumber(usage.output);
	}

	return totals;
}

function estimateTextTokens(text: string): number {
	const trimmed = text.trim();
	if (!trimmed) return 0;
	return Math.max(1, Math.floor(trimmed.length / INPUT_CHARS_PER_TOKEN));
}

function digitCells(value: number, width: number): string[] {
	const digits = Math.max(0, Math.floor(value)).toString().padStart(width, "0").split("");
	const groups: string[] = [];

	for (let i = 0; i < digits.length; i++) {
		groups.push(`[${digits[i]}]`);
		const remaining = digits.length - i - 1;
		if (remaining > 0 && remaining % 3 === 0) groups.push(" ");
	}

	return groups;
}

function renderOdometer(label: string, value: number, target: number, color: (s: string) => string, dim: (s: string) => string): string {
	const shown = Math.round(value);
	const width = Math.max(6, Math.max(shown, target).toString().length);
	const cells = digitCells(shown, width).join("");
	const rolling = shown !== target ? dim(" ↗") : "  ";
	return `${color(label.padEnd(4))} ${cells}${rolling}`;
}

function renderWidget(ctx: ExtensionContext, state: AnimationState): string[] {
	const theme = ctx.ui.theme;
	const input = state.displayed.input;
	const output = state.displayed.output;
	const targetInput = Math.round(state.target.input);
	const targetOutput = Math.round(state.target.output);
	const total = Math.round(input + output);

	return [
		theme.fg("dim", "╭─ token odometers ─────────────────"),
		renderOdometer("IN", input, targetInput, (s) => theme.fg("accent", s), (s) => theme.fg("dim", s)),
		renderOdometer("OUT", output, targetOutput, (s) => theme.fg("success", s), (s) => theme.fg("dim", s)),
		theme.fg("dim", `total ${total.toLocaleString()} tokens in current branch`),
		theme.fg("dim", "╰───────────────────────────────────"),
	];
}

function paint(state: AnimationState): void {
	if (!state.ctx || !state.visible || !state.ctx.hasUI) return;
	state.ctx.ui.setWidget(WIDGET_ID, renderWidget(state.ctx, state), { placement: "aboveEditor" });
}

function stopTimer(state: AnimationState): void {
	if (state.timer) {
		clearInterval(state.timer);
		state.timer = undefined;
	}
}

function tick(state: AnimationState): void {
	const elapsed = Date.now() - state.startedAt;
	const t = easeOutCubic(elapsed / DURATION_MS);

	state.displayed = {
		input: state.from.input + (state.target.input - state.from.input) * t,
		output: state.from.output + (state.target.output - state.from.output) * t,
	};

	paint(state);

	if (t >= 1) {
		state.displayed = { ...state.target };
		paint(state);
		stopTimer(state);
	}
}

function animateTo(state: AnimationState, totals: Totals): void {
	const next = {
		input: Math.max(0, Math.floor(totals.input)),
		output: Math.max(0, Math.floor(totals.output)),
	};

	if (next.input === state.target.input && next.output === state.target.output) {
		paint(state);
		return;
	}

	state.from = { ...state.displayed };
	state.target = next;
	state.startedAt = Date.now();

	if (!state.timer) {
		state.timer = setInterval(() => tick(state), FRAME_MS);
	}
	tick(state);
}

function liveTarget(state: AnimationState): Totals {
	return {
		input: state.committed.input + state.liveInputEstimate,
		output: state.committed.output + state.liveOutputEstimate,
	};
}

function resetLiveEstimates(state: AnimationState): void {
	state.liveInputEstimate = 0;
	state.liveOutputEstimate = 0;
	state.outputCharRemainder = 0;
}

function refreshFromSession(state: AnimationState, ctx: ExtensionContext, immediate = false): void {
	state.ctx = ctx;
	state.committed = readSessionTotals(ctx);
	resetLiveEstimates(state);

	if (immediate) {
		stopTimer(state);
		state.target = { ...state.committed };
		state.from = { ...state.committed };
		state.displayed = { ...state.committed };
		state.startedAt = Date.now();
		paint(state);
		return;
	}

	animateTo(state, state.committed);
}

function settleWithAssistantMessage(state: AnimationState, ctx: ExtensionContext, message: any): void {
	state.ctx = ctx;
	const branchTotals = readSessionTotals(ctx);
	const messageUsage = usageFromMessage(message);
	state.committed = add(branchTotals, messageUsage);
	resetLiveEstimates(state);
	animateTo(state, state.committed);
}

export default function (pi: ExtensionAPI) {
	const state: AnimationState = {
		visible: true,
		committed: { input: 0, output: 0 },
		target: { input: 0, output: 0 },
		from: { input: 0, output: 0 },
		displayed: { input: 0, output: 0 },
		startedAt: Date.now(),
		liveInputEstimate: 0,
		liveOutputEstimate: 0,
		outputCharRemainder: 0,
	};

	pi.on("session_start", async (_event, ctx) => {
		refreshFromSession(state, ctx, true);
	});

	pi.on("input", async (event, ctx) => {
		// Start the input odometer immediately when the user submits text. The
		// provider's exact input usage is not known yet, so this is a live estimate
		// that is reconciled when the assistant response finishes.
		state.ctx = ctx;
		state.committed = readSessionTotals(ctx);
		state.liveInputEstimate += estimateTextTokens(event.text);
		if (event.images?.length) state.liveInputEstimate += event.images.length * 256;
		animateTo(state, liveTarget(state));
		return { action: "continue" };
	});

	pi.on("message_start", async (event, ctx) => {
		if (event.message.role === "assistant") {
			state.ctx = ctx;
			state.liveOutputEstimate = 0;
			state.outputCharRemainder = 0;
		}
	});

	pi.on("message_update", async (event, ctx) => {
		const streamEvent = event.assistantMessageEvent;

		if (streamEvent.type === "text_delta" || streamEvent.type === "thinking_delta" || streamEvent.type === "toolcall_delta") {
			state.ctx = ctx;
			state.outputCharRemainder += streamEvent.delta.length;
			const inc = Math.floor(state.outputCharRemainder / OUTPUT_CHARS_PER_TOKEN);
			if (inc > 0) {
				state.liveOutputEstimate += inc;
				state.outputCharRemainder -= inc * OUTPUT_CHARS_PER_TOKEN;
				animateTo(state, liveTarget(state));
			}
			return;
		}

		if (streamEvent.type === "done") {
			settleWithAssistantMessage(state, ctx, streamEvent.message);
		} else if (streamEvent.type === "error") {
			settleWithAssistantMessage(state, ctx, streamEvent.error);
		}
	});

	pi.on("message_end", async (event, ctx) => {
		if (event.message.role === "assistant") refreshFromSession(state, ctx);
	});

	pi.on("agent_end", async (_event, ctx) => {
		refreshFromSession(state, ctx);
	});

	pi.on("session_tree", async (_event, ctx) => {
		// Tree navigation can move to a branch with lower totals, so resync exactly.
		refreshFromSession(state, ctx, true);
	});

	pi.on("session_compact", async (_event, ctx) => {
		refreshFromSession(state, ctx, true);
	});

	pi.on("session_shutdown", async (_event, ctx) => {
		stopTimer(state);
		ctx.ui.setWidget(WIDGET_ID, undefined);
		ctx.ui.setWorkingMessage();
	});

	pi.registerCommand("token-odometer", {
		description: "Show, hide, or refresh animated input/output token odometers.",
		handler: async (args, ctx) => {
			const action = args.trim().toLowerCase();

			if (action === "hide" || action === "off") {
				state.visible = false;
				stopTimer(state);
				ctx.ui.setWidget(WIDGET_ID, undefined);
				ctx.ui.setWorkingMessage();
				ctx.ui.notify("Token odometer hidden", "info");
				return;
			}

			if (action === "show" || action === "on" || action === "") {
				state.visible = true;
				refreshFromSession(state, ctx, action === "");
				ctx.ui.notify("Token odometer shown", "info");
				return;
			}

			if (action === "refresh") {
				state.visible = true;
				refreshFromSession(state, ctx, true);
				ctx.ui.notify("Token odometer refreshed", "info");
				return;
			}

			ctx.ui.notify("Usage: /token-odometer [show|hide|refresh]", "error");
		},
	});
}