import type { StreamOperation, StreamWireEvent } from "../types";
import type { ChatMessage } from "../types/messages";
import { selectAdapter } from "./protocol";
import type { StreamProtocolAdapter } from "./protocol";
import { streamReducer } from "./reducer";
import type { StreamReducerState } from "./reducer";
import { createAssistantTextMessage } from "../utils/messageParts";
import { chatUnavailableMessage } from "../constants";

export interface StreamingSessionCallbacks {
	appendMessage: (message: ChatMessage) => void;
	replaceMessage: (messageId: string, message: ChatMessage) => void;
}

function createLineParser(): TransformStream<
	Uint8Array,
	Record<string, unknown>
> {
	let buffer = "";
	const decoder = new TextDecoder();

	return new TransformStream({
		transform(chunk, controller) {
			buffer += decoder.decode(chunk, { stream: true });

			const lines = buffer.split("\n");
			buffer = lines.pop() ?? "";

			for (const line of lines) {
				if (line.trim()) {
					try {
						const parsed = JSON.parse(line) as Record<string, unknown>;
						controller.enqueue(parsed);
					} catch {
						controller.enqueue({
							type: "item",
							content: line,
						});
					}
				}
			}
		},

		flush(controller) {
			if (buffer.trim()) {
				try {
					const parsed = JSON.parse(buffer) as Record<string, unknown>;
					controller.enqueue(parsed);
				} catch {
					controller.enqueue({
						type: "item",
						content: buffer,
					});
				}
			}
		},
	});
}

export async function processStreamingResponse(
	body: ReadableStream<Uint8Array>,
	callbacks: StreamingSessionCallbacks,
): Promise<void> {
	const reader = body.pipeThrough(createLineParser()).getReader();
	let adapter: StreamProtocolAdapter | null = null;
	let reducerState: StreamReducerState = { message: null };
	let messageAppended = false;
	let hasReceivedContent = false;

	function applyOperation(op: StreamOperation): void {
		if (
			op.type === "textDelta" ||
			op.type === "productsUpsert" ||
			op.type === "sourceUpsert" ||
			op.type === "statusUpsert" ||
			op.type === "messageError"
		) {
			hasReceivedContent = true;
		}

		reducerState = streamReducer(reducerState, op);
		if (!reducerState.message) return;

		// Defer appending the assistant bubble until the message has
		// visible content (status, text, products, sources, or error).
		if (!hasReceivedContent) return;

		if (messageAppended) {
			callbacks.replaceMessage(
				reducerState.message.id,
				reducerState.message,
			);
		} else {
			callbacks.appendMessage(reducerState.message);
			messageAppended = true;
		}

	}

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

			if (!adapter) {
				const selected = selectAdapter(value as StreamWireEvent);
				if (selected) {
					adapter = selected;
				}
			}

			if (!adapter) continue;

			const ops = adapter.adapt(value as StreamWireEvent);
			for (const op of ops) {
				applyOperation(op);
			}
		}

		if (adapter) {
			const finishOps = adapter.finish();
			for (const op of finishOps) {
				applyOperation(op);
			}
		}
	} catch {
		if (reducerState.message) {
			const errorState = streamReducer(reducerState, {
				type: "messageError",
				message: chatUnavailableMessage,
			});
			if (errorState.message) {
				if (messageAppended) {
					callbacks.replaceMessage(
						errorState.message.id,
						errorState.message,
					);
				} else {
					callbacks.appendMessage(errorState.message);
				}
			}
		} else {
			const errorMsg = createAssistantTextMessage(
				chatUnavailableMessage,
			);
			callbacks.appendMessage(errorMsg);
		}

		console.error("Chat stream failed");
		return;
	} finally {
		reader.releaseLock();
	}

	if (!hasReceivedContent) {
		const emptyMsg = createAssistantTextMessage(
			"[No response received. This could happen if streaming is enabled in the trigger but disabled in agent node(s)]",
		);
		if (messageAppended && reducerState.message) {
			callbacks.replaceMessage(reducerState.message.id, emptyMsg);
		} else {
			callbacks.appendMessage(emptyMsg);
		}
	}
}