import type { Duplex } from "node:stream";

const FRAME_PREFIX_BYTES = 8;

export interface PrivateFrameLimits {
	maxHeaderBytes: number;
	maxPayloadBytes: number;
}

export const DEFAULT_PRIVATE_FRAME_LIMITS: PrivateFrameLimits = {
	maxHeaderBytes: 1024 * 1024,
	maxPayloadBytes: 1024 * 1024 * 1024,
};

export interface PrivateFrame<THeader extends object> {
	header: THeader;
	payload: Buffer;
}

export type PrivateFrameHeaderValidator<THeader extends object> = (value: unknown) => value is THeader;

function assertFrameLength(name: string, value: number, maximum: number): void {
	if (!Number.isSafeInteger(value) || value < 0 || value > maximum) {
		throw new Error(`Invalid private frame ${name}: ${value}`);
	}
}

function isObjectHeader(value: unknown): value is object {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

export function encodePrivateFrame<THeader extends object>(
	header: THeader,
	payload: Uint8Array = Buffer.alloc(0),
	limits: PrivateFrameLimits = DEFAULT_PRIVATE_FRAME_LIMITS,
): Buffer {
	const headerBuffer = Buffer.from(JSON.stringify(header), "utf8");
	const payloadBuffer = Buffer.isBuffer(payload) ? payload : Buffer.from(payload);
	assertFrameLength("header length", headerBuffer.length, limits.maxHeaderBytes);
	assertFrameLength("payload length", payloadBuffer.length, limits.maxPayloadBytes);
	if (headerBuffer.length === 0) {
		throw new Error("Private frame header cannot be empty");
	}

	const frame = Buffer.allocUnsafe(FRAME_PREFIX_BYTES + headerBuffer.length + payloadBuffer.length);
	frame.writeUInt32BE(headerBuffer.length, 0);
	frame.writeUInt32BE(payloadBuffer.length, 4);
	headerBuffer.copy(frame, FRAME_PREFIX_BYTES);
	payloadBuffer.copy(frame, FRAME_PREFIX_BYTES + headerBuffer.length);
	return frame;
}

export class PrivateFrameDecoder<THeader extends object> {
	// Incoming chunks are kept as segments and only joined when a whole frame has
	// arrived. Concatenating into one accumulator on every chunk is O(n^2) in the
	// frame size, which a multi-megabyte response split across socket reads feels
	// as seconds: the session tree of a long transcript, for example.
	private segments: Buffer<ArrayBufferLike>[] = [];
	private pendingBytes = 0;

	constructor(
		private readonly validateHeader: PrivateFrameHeaderValidator<THeader>,
		private readonly limits: PrivateFrameLimits = DEFAULT_PRIVATE_FRAME_LIMITS,
	) {}

	get bufferedBytes(): number {
		return this.pendingBytes;
	}

	push(chunk: Uint8Array): PrivateFrame<THeader>[] {
		if (chunk.length > 0) {
			const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk);
			this.segments.push(buffer);
			this.pendingBytes += buffer.length;
		}

		const frames: PrivateFrame<THeader>[] = [];
		while (this.pendingBytes >= FRAME_PREFIX_BYTES) {
			const prefix = this.peek(FRAME_PREFIX_BYTES);
			const headerLength = prefix.readUInt32BE(0);
			const payloadLength = prefix.readUInt32BE(4);
			assertFrameLength("header length", headerLength, this.limits.maxHeaderBytes);
			assertFrameLength("payload length", payloadLength, this.limits.maxPayloadBytes);
			if (headerLength === 0) {
				throw new Error("Private frame header cannot be empty");
			}

			const frameLength = FRAME_PREFIX_BYTES + headerLength + payloadLength;
			if (this.pendingBytes < frameLength) {
				break;
			}

			// Join the frame's segments exactly once (a whole frame that arrived in
			// one chunk is taken as-is).
			const frame = this.take(frameLength);
			const headerStart = FRAME_PREFIX_BYTES;
			const payloadStart = headerStart + headerLength;
			let decoded: unknown;
			try {
				decoded = JSON.parse(frame.toString("utf8", headerStart, payloadStart));
			} catch (error) {
				throw new Error(
					`Invalid private frame header JSON: ${error instanceof Error ? error.message : String(error)}`,
				);
			}
			if (!isObjectHeader(decoded) || !this.validateHeader(decoded)) {
				throw new Error("Invalid private frame routing header");
			}

			frames.push({
				header: decoded,
				payload: Buffer.from(frame.subarray(payloadStart, payloadStart + payloadLength)),
			});
		}
		return frames;
	}

	/** Copy of the next `count` buffered bytes, without consuming them. */
	private peek(count: number): Buffer {
		const first = this.segments[0];
		if (first !== undefined && first.length >= count) {
			return first.subarray(0, count);
		}
		const head: Buffer[] = [];
		let remaining = count;
		for (const segment of this.segments) {
			const take = Math.min(segment.length, remaining);
			head.push(take === segment.length ? segment : segment.subarray(0, take));
			remaining -= take;
			if (remaining === 0) break;
		}
		return Buffer.concat(head, count);
	}

	/** Remove and return the next `count` buffered bytes as one contiguous buffer. */
	private take(count: number): Buffer {
		const head: Buffer[] = [];
		let remaining = count;
		let consumed = 0;
		while (remaining > 0) {
			const segment = this.segments[consumed]!;
			if (segment.length <= remaining) {
				head.push(segment);
				remaining -= segment.length;
				consumed++;
			} else {
				head.push(segment.subarray(0, remaining));
				// Keep the unconsumed tail as a copy so the caller's chunk is released.
				this.segments[consumed] = Buffer.from(segment.subarray(remaining));
				remaining = 0;
			}
		}
		this.segments = this.segments.slice(consumed);
		this.pendingBytes -= count;
		return head.length === 1 ? head[0]! : Buffer.concat(head, count);
	}

	finish(): void {
		if (this.pendingBytes !== 0) {
			throw new Error(`Private frame channel ended with ${this.pendingBytes} incomplete bytes`);
		}
	}
}

export type PrivateFrameListener<THeader extends object> = (frame: PrivateFrame<THeader>) => void;

export class PrivateFramedChannel<THeader extends object> {
	private readonly decoder: PrivateFrameDecoder<THeader>;
	private readonly listeners = new Set<PrivateFrameListener<THeader>>();
	private closed = false;

	constructor(
		private readonly stream: Duplex,
		validateHeader: PrivateFrameHeaderValidator<THeader>,
		private readonly limits: PrivateFrameLimits = DEFAULT_PRIVATE_FRAME_LIMITS,
	) {
		this.decoder = new PrivateFrameDecoder(validateHeader, limits);
		stream.on("data", this.handleData);
		stream.on("end", this.handleEnd);
		stream.on("close", this.handleClose);
	}

	onFrame(listener: PrivateFrameListener<THeader>): () => void {
		this.listeners.add(listener);
		return () => this.listeners.delete(listener);
	}

	async send(header: THeader, payload?: Uint8Array): Promise<void> {
		if (this.closed || this.stream.destroyed) {
			throw new Error("Private frame channel is closed");
		}
		const frame = encodePrivateFrame(header, payload, this.limits);
		await new Promise<void>((resolve, reject) => {
			this.stream.write(frame, (error?: Error | null) => {
				if (error) {
					reject(error);
				} else {
					resolve();
				}
			});
		});
	}

	close(): void {
		if (this.closed) {
			return;
		}
		this.closed = true;
		this.detach();
		this.stream.end();
	}

	private readonly handleData = (chunk: Buffer): void => {
		try {
			for (const frame of this.decoder.push(chunk)) {
				for (const listener of this.listeners) {
					listener(frame);
				}
			}
		} catch (error) {
			this.stream.destroy(error instanceof Error ? error : new Error(String(error)));
		}
	};

	private readonly handleEnd = (): void => {
		try {
			this.decoder.finish();
		} catch (error) {
			this.stream.destroy(error instanceof Error ? error : new Error(String(error)));
		}
	};

	private readonly handleClose = (): void => {
		this.closed = true;
		this.detach();
	};

	private detach(): void {
		this.stream.off("data", this.handleData);
		this.stream.off("end", this.handleEnd);
		this.stream.off("close", this.handleClose);
		this.listeners.clear();
	}
}