import { describe, expect, it } from "vitest";
import { calculateLiveTps, calculateReliableTps, setupTpsTracker, type TpsTiming } from "./tps.ts";

function timing(overrides: Partial<TpsTiming> = {}): TpsTiming {
	return {
		messageStartMs: 0,
		lastUpdateMs: 350,
		firstTokenMs: 100,
		updateCount: 5,
		firstStreamUpdateMs: 150,
		lastStreamUpdateMs: 350,
		stallMs: 0,
		generationMs: 400,
		...overrides,
	};
}

describe("calculateLiveTps", () => {
	it("waits for a usable sample and rejects implausible rates", () => {
		expect(calculateLiveTps(60, 0.2)).toBeNull();
		expect(calculateLiveTps(100, 200)).toBe(500);
		expect(calculateLiveTps(10_001, 1_000)).toBeNull();
	});
});

describe("setupTpsTracker with anthropic-style usage", () => {
	it("ignores the tiny usage.output seeded at message_start while streaming", async () => {
		const handlers = new Map<string, (event: any, ctx: any) => Promise<void>>();
		const statuses: string[] = [];
		const ctx = { ui: { setStatus: (_k: string, v: string) => statuses.push(v), notify: () => {}, theme: { fg: (_c: string, t: string) => t } } };
		setupTpsTracker({ on: (name: string, fn: any) => handlers.set(name, fn) } as any);

		let now = 0;
		const realNow = performance.now;
		performance.now = () => now;
		try {
			const message = { role: "assistant", provider: "anthropic", model: "claude", usage: { output: 1 } };
			await handlers.get("agent_start")!({}, ctx);
			await handlers.get("message_start")!({ message }, ctx);
			for (let i = 0; i < 20; i++) {
				now += 50;
				await handlers.get("message_update")!({ message, assistantMessageEvent: { type: "text_delta", delta: "x".repeat(40) } }, ctx);
			}
			// 20 deltas * 10 est. tokens over ~1s => ~200 tok/s, not ~1 tok/s
			expect(statuses.at(-1)).toBe("211 tok/s");

			message.usage.output = 200;
			await handlers.get("message_end")!({ message }, ctx);
			await handlers.get("agent_end")!({}, ctx);
			expect(statuses.at(-1)).toMatch(/^done [1-9]\d* t\/s \(main [1-9]\d* t\/s\)$/);
		} finally {
			performance.now = realNow;
		}
	});
});

describe("calculateReliableTps", () => {
	it("uses active stream time for sufficiently sampled output", () => {
		expect(calculateReliableTps(100, timing())).toEqual({
			tps: 500,
			effectiveMs: 200,
			isPrimary: true,
		});
	});

	it("falls back around stalls and rejects implausible samples", () => {
		expect(
			calculateReliableTps(100, timing({ updateCount: 2, firstStreamUpdateMs: 150, lastStreamUpdateMs: 250, generationMs: 500 })),
		).toEqual({ tps: 200, effectiveMs: 500, isPrimary: false });
		expect(calculateReliableTps(100, timing({ stallMs: 100 }))).toEqual({
			tps: 333.3,
			effectiveMs: 300,
			isPrimary: false,
		});
		expect(calculateReliableTps(10_001, timing())).toBeNull();
	});
});