import type { AssistantMessage, AssistantMessageEvent, UserMessage } from "@earendil-works/pi-ai";
import { describe, expect, test } from "vitest";
import { AgentActivityTracker, formatTokenCount } from "../src/modes/interactive/agent-activity.js";

function createAssistantMessage(outputTokens = 0): AssistantMessage {
	return {
		role: "assistant",
		content: [],
		api: "anthropic-messages",
		provider: "anthropic",
		model: "test-model",
		usage: {
			input: 0,
			output: outputTokens,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: outputTokens,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		},
		stopReason: "stop",
		timestamp: 0,
	};
}

function createUserMessage(): UserMessage {
	return { role: "user", content: "hello", timestamp: 0 };
}

function update(message: AssistantMessage, assistantMessageEvent: AssistantMessageEvent) {
	return { type: "message_update" as const, message, assistantMessageEvent };
}

describe("AgentActivityTracker", () => {
	test("starts waiting with no tokens", () => {
		const tracker = new AgentActivityTracker();
		expect(tracker.getStatus()).toEqual({ activity: "waiting", direction: "up", tokens: 0 });
	});

	test("derives activity from streaming delta types", () => {
		const tracker = new AgentActivityTracker();
		const message = createAssistantMessage();
		tracker.handleEvent({ type: "agent_start" });
		tracker.handleEvent({ type: "message_start", message });

		tracker.handleEvent(update(message, { type: "thinking_delta", contentIndex: 0, delta: "hmm", partial: message }));
		expect(tracker.getStatus().activity).toBe("thinking");
		expect(tracker.getStatus().direction).toBe("down");

		tracker.handleEvent(update(message, { type: "text_delta", contentIndex: 1, delta: "hi", partial: message }));
		expect(tracker.getStatus().activity).toBe("writing");

		tracker.handleEvent(update(message, { type: "toolcall_start", contentIndex: 2, partial: message }));
		expect(tracker.getStatus().activity).toBe("writing-code");
		expect(tracker.getStatus().direction).toBe("down");
	});

	test("uses streamed usage tokens when the provider reports them", () => {
		const tracker = new AgentActivityTracker();
		const message = createAssistantMessage(1234);
		tracker.handleEvent({ type: "message_start", message });
		tracker.handleEvent(update(message, { type: "text_delta", contentIndex: 0, delta: "hi", partial: message }));
		expect(tracker.getStatus().tokens).toBe(1234);
	});

	test("falls back to a character estimate when usage is not streamed", () => {
		const tracker = new AgentActivityTracker();
		const message = createAssistantMessage(0);
		tracker.handleEvent({ type: "message_start", message });
		tracker.handleEvent(
			update(message, { type: "text_delta", contentIndex: 0, delta: "a".repeat(40), partial: message }),
		);
		expect(tracker.getStatus().tokens).toBe(10);
	});

	test("updates live from the character estimate when streamed usage is stale", () => {
		const tracker = new AgentActivityTracker();
		const message = createAssistantMessage(1);
		tracker.handleEvent({ type: "message_start", message });
		tracker.handleEvent(
			update(message, { type: "text_delta", contentIndex: 0, delta: "a".repeat(400), partial: message }),
		);
		expect(tracker.getStatus().tokens).toBe(100);
	});

	test("never decreases when authoritative usage arrives at message end", () => {
		const tracker = new AgentActivityTracker();
		const message = createAssistantMessage(0);
		tracker.handleEvent({ type: "message_start", message });
		tracker.handleEvent(
			update(message, { type: "text_delta", contentIndex: 0, delta: "a".repeat(400), partial: message }),
		);
		expect(tracker.getStatus().tokens).toBe(100);
		tracker.handleEvent({ type: "message_end", message: createAssistantMessage(80) });
		expect(tracker.getStatus().tokens).toBe(100);
	});

	test("accumulates tokens across turns and flips direction while executing tools", () => {
		const tracker = new AgentActivityTracker();
		const first = createAssistantMessage(500);
		tracker.handleEvent({ type: "message_start", message: first });
		tracker.handleEvent(update(first, { type: "toolcall_delta", contentIndex: 0, delta: "code", partial: first }));
		tracker.handleEvent({ type: "message_end", message: first });
		expect(tracker.getStatus()).toEqual({ activity: "waiting", direction: "up", tokens: 500 });

		tracker.handleEvent({ type: "tool_execution_start", toolCallId: "t1", toolName: "ipython", args: {} });
		expect(tracker.getStatus()).toEqual({ activity: "executing", direction: "up", tokens: 500 });

		tracker.handleEvent({
			type: "tool_execution_end",
			toolCallId: "t1",
			toolName: "ipython",
			result: {},
			isError: false,
		});
		expect(tracker.getStatus().activity).toBe("waiting");

		const second = createAssistantMessage(250);
		tracker.handleEvent({ type: "message_start", message: second });
		tracker.handleEvent(update(second, { type: "text_delta", contentIndex: 0, delta: "done", partial: second }));
		expect(tracker.getStatus().tokens).toBe(750);
	});

	test("stays executing while a parallel tool call is still running", () => {
		const tracker = new AgentActivityTracker();
		tracker.handleEvent({ type: "tool_execution_start", toolCallId: "t1", toolName: "ipython", args: {} });
		tracker.handleEvent({ type: "tool_execution_start", toolCallId: "t2", toolName: "ipython", args: {} });
		tracker.handleEvent({
			type: "tool_execution_end",
			toolCallId: "t1",
			toolName: "ipython",
			result: {},
			isError: false,
		});
		expect(tracker.getStatus().activity).toBe("executing");
		tracker.handleEvent({
			type: "tool_execution_end",
			toolCallId: "t2",
			toolName: "ipython",
			result: {},
			isError: false,
		});
		expect(tracker.getStatus().activity).toBe("waiting");
	});

	test("resets token count on the next user message", () => {
		const tracker = new AgentActivityTracker();
		const message = createAssistantMessage(500);
		tracker.handleEvent({ type: "message_start", message });
		tracker.handleEvent({ type: "message_end", message });
		expect(tracker.getStatus().tokens).toBe(500);

		tracker.handleEvent({ type: "message_start", message: createUserMessage() });
		expect(tracker.getStatus()).toEqual({ activity: "waiting", direction: "up", tokens: 0 });
	});
});

describe("formatTokenCount", () => {
	test("formats counts below 1k as-is", () => {
		expect(formatTokenCount(0)).toBe("0");
		expect(formatTokenCount(999)).toBe("999");
	});

	test("formats small thousands with one decimal", () => {
		expect(formatTokenCount(1000)).toBe("1.0k");
		expect(formatTokenCount(9940)).toBe("9.9k");
	});

	test("drops the decimal at 10k and above", () => {
		expect(formatTokenCount(12340)).toBe("12k");
		expect(formatTokenCount(250400)).toBe("250k");
	});

	test("formats millions", () => {
		expect(formatTokenCount(1500000)).toBe("1.5M");
		expect(formatTokenCount(12000000)).toBe("12M");
	});
});