import { libx } from 'libx.js/build/bundles/essentials.js';

/**
 * BlackoutManager - Redacts content within <|...|> sections during streaming
 * Useful for protecting sensitive information in prompts and responses
 */
export class BlackoutManager {
	private enabled: boolean;
	private redactionMarker: string = '[REDACTED]';

	constructor(enabled: boolean = false) {
		this.enabled = enabled;
	}

	/**
	 * Process a text chunk and redact any content within <|...|> markers
	 * Handles partial markers across chunk boundaries
	 */
	public processChunk(chunk: string, state: BlackoutState): { output: string; state: BlackoutState; } {
		if (!this.enabled) {
			return { output: chunk, state };
		}

		let output = '';
		let currentState = state.inRedactedSection;
		let buffer = state.buffer + chunk;
		let i = 0;

		while (i < buffer.length) {
			if (!currentState) {
				// Look for opening marker <|
				const openIndex = buffer.indexOf('<|', i);

				if (openIndex === -1) {
					// No opening marker found, output rest and save potential partial marker
					if (buffer.length - i > 1) {
						// Keep last character in case it's part of '<|'
						output += buffer.slice(i, -1);
						state.buffer = buffer.slice(-1);
					} else {
						state.buffer = buffer.slice(i);
					}
					break;
				}

				// Output everything before the marker
				output += buffer.slice(i, openIndex);
				i = openIndex + 2;
				currentState = true;
			} else {
				// Look for closing marker |>
				const closeIndex = buffer.indexOf('|>', i);

				if (closeIndex === -1) {
					// No closing marker found, buffer the rest
					state.buffer = buffer.slice(i);
					break;
				}

				// Skip redacted content and output marker
				output += this.redactionMarker;
				i = closeIndex + 2;
				currentState = false;
				state.buffer = '';
			}
		}

		state.inRedactedSection = currentState;

		libx.log.v('BlackoutManager: processed chunk', {
			inputLength: chunk.length,
			outputLength: output.length,
			inRedactedSection: currentState,
			bufferSize: state.buffer.length
		});

		return { output, state };
	}

	/**
	 * Process complete text and redact all <|...|> sections
	 */
	public processText(text: string): string {
		if (!this.enabled) return text;

		let result = text;
		const regex = /<\|.*?\|>/gs;
		result = result.replace(regex, this.redactionMarker);

		libx.log.v('BlackoutManager: processed text', {
			inputLength: text.length,
			outputLength: result.length,
			redactionsCount: (text.match(regex) || []).length
		});

		return result;
	}

	/**
	 * Create a transform stream that redacts content on the fly
	 */
	public createRedactionStream(): TransformStream<Uint8Array, Uint8Array> {
		const encoder = new TextEncoder();
		const decoder = new TextDecoder();
		const state: BlackoutState = { inRedactedSection: false, buffer: '' };
		const self = this;

		return new TransformStream({
			transform(chunk, controller) {
				const text = decoder.decode(chunk, { stream: true });
				const { output } = self.processChunk(text, state);

				if (output) {
					controller.enqueue(encoder.encode(output));
				}
			},
			flush(controller) {
				// Flush any remaining buffer
				if (state.buffer) {
					controller.enqueue(encoder.encode(state.buffer));
				}
			}
		});
	}
}

export interface BlackoutState {
	inRedactedSection: boolean;
	buffer: string;
}

/**
 * Apply blackout to a stream response
 */
export async function applyBlackoutToStream(
	stream: ReadableStream,
	enabled: boolean
): Promise<ReadableStream> {
	if (!enabled) return stream;

	const blackout = new BlackoutManager(true);
	return stream.pipeThrough(blackout.createRedactionStream());
}