import assert from "node:assert/strict";
import { describe, it } from "vitest";
import { parseAnthropicSSE } from "./event-stream.js";

function createMockResponse(chunks: string[]): Response {
	let index = 0;
	const encoder = new TextEncoder();
	const stream = new ReadableStream<Uint8Array>({
		pull(controller) {
			if (index < chunks.length) {
				controller.enqueue(encoder.encode(chunks[index++]));
			} else {
				controller.close();
			}
		},
	});
	return new Response(stream);
}

describe("parseAnthropicSSE", () => {
	it("yields parsed JSON for known Anthropic events", async () => {
		const sse =
			"event: message_start\n" +
			'data: {"type":"message_start","message":{"id":"msg_1","role":"assistant","content":[],"model":"claude-3","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}}\n' +
			"\n" +
			"event: content_block_start\n" +
			'data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n' +
			"\n" +
			"event: content_block_delta\n" +
			'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}\n' +
			"\n" +
			"event: content_block_stop\n" +
			'data: {"type":"content_block_stop","index":0}\n' +
			"\n" +
			"event: message_delta\n" +
			'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":10,"output_tokens":1,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}\n' +
			"\n" +
			"event: message_stop\n" +
			'data: {"type":"message_stop"}\n' +
			"\n";

		const response = createMockResponse([sse]);
		const events: unknown[] = [];
		for await (const event of parseAnthropicSSE(response)) {
			events.push(event);
		}

		assert.equal(events.length, 6);
		assert.equal((events[0] as any).type, "message_start");
		assert.equal((events[1] as any).type, "content_block_start");
		assert.equal((events[2] as any).type, "content_block_delta");
		assert.equal((events[3] as any).type, "content_block_stop");
		assert.equal((events[4] as any).type, "message_delta");
		assert.equal((events[5] as any).type, "message_stop");
	});

	it("silently drops unknown events (e.g. OpenAI-style done)", async () => {
		const sse =
			"event: message_start\n" +
			'data: {"type":"message_start","message":{"id":"msg_1","role":"assistant","content":[],"model":"claude-3","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}}\n' +
			"\n" +
			"event: done\n" +
			"data: [DONE]\n" +
			"\n" +
			"event: content_block_start\n" +
			'data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n' +
			"\n";

		const response = createMockResponse([sse]);
		const events: unknown[] = [];
		for await (const event of parseAnthropicSSE(response)) {
			events.push(event);
		}

		assert.equal(events.length, 2);
		assert.equal((events[0] as any).type, "message_start");
		assert.equal((events[1] as any).type, "content_block_start");
	});

	it("ignores ping events", async () => {
		const sse =
			"event: message_start\n" +
			'data: {"type":"message_start","message":{"id":"msg_1","role":"assistant","content":[],"model":"claude-3","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}}\n' +
			"\n" +
			"event: ping\n" +
			"data: {}\n" +
			"\n" +
			"event: message_stop\n" +
			'data: {"type":"message_stop"}\n' +
			"\n";

		const response = createMockResponse([sse]);
		const events: unknown[] = [];
		for await (const event of parseAnthropicSSE(response)) {
			events.push(event);
		}

		assert.equal(events.length, 2);
		assert.equal((events[0] as any).type, "message_start");
		assert.equal((events[1] as any).type, "message_stop");
	});

	it("handles chunked SSE data across multiple reads", async () => {
		const chunks = [
			"event: message_start\n",
			'data: {"type":"message_start","message":{"id":"msg_1","role":"assistant","content":[],"model":"claude-3","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}}\n\n',
			"event: message_stop\n",
			'data: {"type":"message_stop"}\n\n',
		];

		const response = createMockResponse(chunks);
		const events: unknown[] = [];
		for await (const event of parseAnthropicSSE(response)) {
			events.push(event);
		}

		assert.equal(events.length, 2);
		assert.equal((events[0] as any).type, "message_start");
		assert.equal((events[1] as any).type, "message_stop");
	});

	it("handles comment lines", async () => {
		const sse =
			": comment line\n" +
			"event: message_start\n" +
			'data: {"type":"message_start","message":{"id":"msg_1","role":"assistant","content":[],"model":"claude-3","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0,"cache_creation_input_tokens":0,"cache_read_input_tokens":0}}}\n' +
			"\n";

		const response = createMockResponse([sse]);
		const events: unknown[] = [];
		for await (const event of parseAnthropicSSE(response)) {
			events.push(event);
		}

		assert.equal(events.length, 1);
		assert.equal((events[0] as any).type, "message_start");
	});
});
