import type { AgentEvent, StepId } from "@dispatch/kernel";
import { describe, expect, it } from "vitest";
import { createMetricsAccumulator } from "./metrics.js";

function stepId(id: string): StepId {
	return id as StepId;
}

describe("createMetricsAccumulator", () => {
	it("builds a TurnMetrics from a single-step turn", () => {
		const acc = createMetricsAccumulator();

		const usageEvent: AgentEvent = {
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		};
		const stepCompleteEvent: AgentEvent = {
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			ttftMs: 50,
			decodeMs: 150,
			genTotalMs: 200,
		};
		const doneEvent: AgentEvent = {
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
			durationMs: 500,
			usage: { inputTokens: 10, outputTokens: 5 },
		};

		acc.ingest(usageEvent);
		acc.ingest(stepCompleteEvent);
		acc.ingest(doneEvent);

		const tm = acc.build("t1");
		expect(tm.turnId).toBe("t1");
		expect(tm.usage.inputTokens).toBe(10);
		expect(tm.usage.outputTokens).toBe(5);
		expect(tm.durationMs).toBe(500);
		expect(tm.steps).toHaveLength(1);
		expect(tm.steps[0]?.stepId).toBe(stepId("t1#0"));
		expect(tm.steps[0]?.usage.inputTokens).toBe(10);
		expect(tm.steps[0]?.usage.outputTokens).toBe(5);
		expect(tm.steps[0]?.ttftMs).toBe(50);
		expect(tm.steps[0]?.decodeMs).toBe(150);
		expect(tm.steps[0]?.genTotalMs).toBe(200);
	});

	it("aggregates multi-step usage and preserves step order", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		});
		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			genTotalMs: 100,
		});
		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#1"),
			usage: { inputTokens: 20, outputTokens: 10 },
		});
		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#1"),
			genTotalMs: 150,
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
			usage: { inputTokens: 30, outputTokens: 15 },
		});

		const tm = acc.build("t1");
		expect(tm.steps).toHaveLength(2);
		expect(tm.steps[0]?.stepId).toBe(stepId("t1#0"));
		expect(tm.steps[0]?.usage.inputTokens).toBe(10);
		expect(tm.steps[1]?.stepId).toBe(stepId("t1#1"));
		expect(tm.steps[1]?.usage.inputTokens).toBe(20);
		expect(tm.usage.inputTokens).toBe(30);
		expect(tm.usage.outputTokens).toBe(15);
	});

	it("joins per-step timing and usage by stepId", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			ttftMs: 50,
			decodeMs: 100,
			genTotalMs: 150,
		});
		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		});

		const tm = acc.build("t1");
		expect(tm.steps).toHaveLength(1);
		expect(tm.steps[0]?.usage.inputTokens).toBe(10);
		expect(tm.steps[0]?.ttftMs).toBe(50);
		expect(tm.steps[0]?.decodeMs).toBe(100);
		expect(tm.steps[0]?.genTotalMs).toBe(150);
	});

	it("turn-level usage comes from done event, not step sum", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		});
		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
			usage: { inputTokens: 100, outputTokens: 50 },
		});

		const tm = acc.build("t1");
		expect(tm.usage.inputTokens).toBe(100);
		expect(tm.usage.outputTokens).toBe(50);
	});

	it("falls back to summing step usage when done.usage is absent", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		});
		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
		});

		const tm = acc.build("t1");
		expect(tm.usage.inputTokens).toBe(10);
		expect(tm.usage.outputTokens).toBe(5);
	});

	it("tolerates missing usage on a step (timing only)", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			genTotalMs: 200,
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
		});

		const tm = acc.build("t1");
		expect(tm.steps).toHaveLength(1);
		expect(tm.steps[0]?.usage.inputTokens).toBe(0);
		expect(tm.steps[0]?.usage.outputTokens).toBe(0);
		expect(tm.steps[0]?.genTotalMs).toBe(200);
	});

	it("tolerates missing timing on a step (usage only)", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
		});

		const tm = acc.build("t1");
		expect(tm.steps).toHaveLength(1);
		expect(tm.steps[0]?.usage.inputTokens).toBe(10);
		expect(tm.steps[0]?.ttftMs).toBeUndefined();
		expect(tm.steps[0]?.decodeMs).toBeUndefined();
		expect(tm.steps[0]?.genTotalMs).toBeUndefined();
	});

	it("reset clears all accumulated state", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
			usage: { inputTokens: 10, outputTokens: 5 },
		});

		acc.reset();

		const tm = acc.build("t2");
		expect(tm.steps).toHaveLength(0);
		expect(tm.usage.inputTokens).toBe(0);
		expect(tm.usage.outputTokens).toBe(0);
	});

	it("contextSize equals inputTokens + outputTokens for a single-step turn", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		});
		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
			usage: { inputTokens: 10, outputTokens: 5 },
		});

		const tm = acc.build("t1");
		expect(tm.contextSize).toBe(15);
	});

	it("contextSize equals ONLY the last step's inputTokens + outputTokens for a multi-step turn", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			usage: { inputTokens: 10, outputTokens: 5 },
		});
		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
		});
		acc.ingest({
			type: "usage",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#1"),
			usage: { inputTokens: 20, outputTokens: 10 },
		});
		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#1"),
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
			usage: { inputTokens: 100, outputTokens: 50 },
		});

		const tm = acc.build("t1");
		expect(tm.contextSize).toBe(30);
		expect(tm.contextSize).not.toBe(tm.usage.inputTokens);
	});

	it("contextSize is undefined when the turn has no steps", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
		});

		const tm = acc.build("t1");
		expect(tm.contextSize).toBeUndefined();
	});

	it("contextSize is undefined when the last step has no usable per-step usage", () => {
		const acc = createMetricsAccumulator();

		acc.ingest({
			type: "step-complete",
			conversationId: "c1",
			turnId: "t1",
			stepId: stepId("t1#0"),
			genTotalMs: 200,
		});
		acc.ingest({
			type: "done",
			conversationId: "c1",
			turnId: "t1",
			reason: "stop",
		});

		const tm = acc.build("t1");
		expect(tm.contextSize).toBeUndefined();
	});
});