import type { AssistantMessage, AssistantMessageEvent } from "../types.js";

/** Known Anthropic SSE event types that we handle. Unknown events are silently dropped. */
const KNOWN_ANTHROPIC_EVENTS = new Set([
	"message_start",
	"message_delta",
	"message_stop",
	"content_block_start",
	"content_block_delta",
	"content_block_stop",
	"ping",
	"error",
]);

/**
 * Parse a raw SSE (Server-Sent Events) stream response into JSON events.
 *
 * Purpose: give us full control over SSE parsing so that non-Anthropic events
 * (e.g. OpenAI-style "done" events injected by proxies) are silently dropped
 * instead of corrupting the stream.
 *
 * Consumer: processAnthropicStream in anthropic-shared.ts.
 */
export async function* parseAnthropicSSE(
	response: Response,
	signal?: AbortSignal,
): AsyncGenerator<unknown, void, unknown> {
	if (!response.body) {
		throw new Error("Attempted to iterate over a response with no body");
	}

	const reader = response.body.getReader();
	const decoder = new TextDecoder();
	let buffer = "";
	let eventName: string | null = null;
	let dataLines: string[] = [];

	function flushEvent(): unknown | undefined {
		if (eventName === null && dataLines.length === 0) {
			return undefined;
		}
		const data = dataLines.join("\n");
		const name = eventName ?? "";
		eventName = null;
		dataLines = [];

		if (name === "ping") {
			return undefined;
		}
		if (name === "error") {
			try {
				return JSON.parse(data);
			} catch {
				return undefined;
			}
		}
		if (!KNOWN_ANTHROPIC_EVENTS.has(name)) {
			// Silently drop unknown events (e.g. OpenAI-style "done" from proxies)
			return undefined;
		}
		try {
			return JSON.parse(data);
		} catch {
			return undefined;
		}
	}

	function processLine(line: string): unknown | undefined {
		const trimmed = line.trim();
		if (trimmed === "") {
			// Empty line means end of an SSE event
			return flushEvent();
		}

		if (trimmed.startsWith(":")) {
			// Comment line, ignore
			return undefined;
		}

		const colonIndex = trimmed.indexOf(":");
		if (colonIndex === -1) return undefined;

		const field = trimmed.slice(0, colonIndex);
		const value = trimmed.slice(colonIndex + 1).trimStart();

		if (field === "event") {
			eventName = value;
		} else if (field === "data") {
			dataLines.push(value);
		}
		return undefined;
	}

	try {
		while (true) {
			if (signal?.aborted) {
				return;
			}

			const { done, value } = await reader.read();
			if (done) break;

			buffer += decoder.decode(value, { stream: true });

			let newlineIndex: number;
			while ((newlineIndex = buffer.indexOf("\n")) !== -1) {
				const line = buffer.slice(0, newlineIndex);
				buffer = buffer.slice(newlineIndex + 1);
				const event = processLine(line);
				if (event !== undefined) {
					yield event;
				}
			}
		}

		// Flush remaining buffer as a final line
		if (buffer.length > 0) {
			const event = processLine(buffer);
			if (event !== undefined) {
				yield event;
			}
		}

		// Flush any pending event
		const event = flushEvent();
		if (event !== undefined) {
			yield event;
		}
	} finally {
		reader.releaseLock();
	}
}

// Generic event stream class for async iteration
export class EventStream<T, R = T> implements AsyncIterable<T> {
	private queue: T[] = [];
	private waiting: ((value: IteratorResult<T>) => void)[] = [];
	private done = false;
	private finalResultPromise: Promise<R>;
	private resolveFinalResult!: (result: R) => void;

	constructor(
		private isComplete: (event: T) => boolean,
		private extractResult: (event: T) => R,
	) {
		this.finalResultPromise = new Promise((resolve) => {
			this.resolveFinalResult = resolve;
		});
	}

	push(event: T): void {
		if (this.done) return;

		if (this.isComplete(event)) {
			this.done = true;
			this.resolveFinalResult(this.extractResult(event));
		}

		// Deliver to waiting consumer or queue it
		const waiter = this.waiting.shift();
		if (waiter) {
			waiter({ value: event, done: false });
		} else {
			this.queue.push(event);
		}
	}

	end(result?: R): void {
		this.done = true;
		if (result !== undefined) {
			this.resolveFinalResult(result);
		}
		// Notify all waiting consumers that we're done
		while (this.waiting.length > 0) {
			const waiter = this.waiting.shift()!;
			waiter({ value: undefined as any, done: true });
		}
	}

	async *[Symbol.asyncIterator](): AsyncIterator<T> {
		while (true) {
			if (this.queue.length > 0) {
				yield this.queue.shift()!;
			} else if (this.done) {
				return;
			} else {
				const result = await new Promise<IteratorResult<T>>((resolve) =>
					this.waiting.push(resolve),
				);
				if (result.done) return;
				yield result.value;
			}
		}
	}

	result(): Promise<R> {
		return this.finalResultPromise;
	}
}

export class AssistantMessageEventStream extends EventStream<
	AssistantMessageEvent,
	AssistantMessage
> {
	constructor() {
		super(
			(event) => event.type === "done" || event.type === "error",
			(event) => {
				if (event.type === "done") {
					return event.message;
				} else if (event.type === "error") {
					return event.error;
				}
				throw new Error("Unexpected event type for final result");
			},
		);
	}
}

/** Factory function for AssistantMessageEventStream (for use by package consumers). */
export function createAssistantMessageEventStream(): AssistantMessageEventStream {
	return new AssistantMessageEventStream();
}
