import type { ProviderToolLiveTracker } from "./provider-tool-live-status";
import type { ProviderToolType } from "./types";

export type JsonRecord = Record<string, unknown>;

const INTERRUPT_DONE_EVENT = "data: [DONE]\n\n";
export const DEFAULT_IMAGE_GENERATION_KEEPALIVE_INTERVAL_MS = 60_000;
const KEEPALIVE_EVENT_TYPE = "response.function_call_arguments.delta";
const KEEPALIVE_ITEM_ID = "openai_provider_tools_keepalive";

const WEB_SEARCH_LIFECYCLE_EVENTS = new Set([
	"response.web_search_call.in_progress",
	"response.web_search_call.searching",
	"response.web_search_call.completed",
]);
export interface RequestObservationPolicy {
	enabledTools?: readonly ProviderToolType[];
	interruptOnImageResult: boolean;
	keepaliveIntervalMs: number | undefined;
	liveTracker?: ProviderToolLiveTracker;
	observeLiveEventsInIterable?: boolean;
}

export interface RequestObservationPolicyBinding {
	currentPolicy(): RequestObservationPolicy | undefined;
	onPolicy(callback: (policy: RequestObservationPolicy) => void): void;
	dispose(): void;
}

type RequestObservationPolicySource = RequestObservationPolicy | RequestObservationPolicyBinding;

function isPolicyBinding(source: RequestObservationPolicySource): source is RequestObservationPolicyBinding {
	return typeof (source as RequestObservationPolicyBinding).currentPolicy === "function"
		&& typeof (source as RequestObservationPolicyBinding).onPolicy === "function";
}

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

export function tryEnqueueChunk(
	controller: ReadableStreamDefaultController<Uint8Array>,
	chunk: Uint8Array,
	markFinished: () => void,
): boolean {
	try {
		controller.enqueue(chunk);
		return true;
	} catch (error) {
		if (isClosedControllerError(error)) {
			markFinished();
			return false;
		}
		throw error;
	}
}

function isClosedControllerError(error: unknown): boolean {
	if (!(error instanceof TypeError)) return false;
	const message = error.message.toLowerCase();
	return message.includes("controller is already closed") || (message.includes("invalid state") && message.includes("closed"));
}

export function wrapOpenAIResponsesStream(body: ReadableStream<Uint8Array>, policySource: RequestObservationPolicySource): ReadableStream<Uint8Array> {
	const reader = body.getReader();
	const decoder = new TextDecoder();
	const encoder = new TextEncoder();
	let policy = isPolicyBinding(policySource) ? policySource.currentPolicy() : policySource;
	let buffer = "";
	let finished = false;
	let keepaliveTimer: ReturnType<typeof setTimeout> | undefined;
	let keepaliveController: ReadableStreamDefaultController<Uint8Array> | undefined;
	let sawLiveWebSearchEvent = false;
	let lastSyntheticKeepaliveAt = Date.now();
	let replayableLiveEvents: unknown[] = [];

	const disposeBinding = () => {
		if (isPolicyBinding(policySource)) policySource.dispose();
	};
	const imageKeepaliveEnabled = () => {
		const enabledTools = policy?.enabledTools;
		return policy !== undefined && (enabledTools === undefined || enabledTools.includes("image_generation"));
	};
	const clearKeepalive = () => {
		if (keepaliveTimer) clearTimeout(keepaliveTimer);
		keepaliveTimer = undefined;
	};
	const finish = () => {
		finished = true;
		clearKeepalive();
		disposeBinding();
	};
	const failAndClearTracker = (error: unknown) => {
		if (!policy) return;
		failTracker(policy.liveTracker, error);
		clearTracker(policy.liveTracker);
	};
	const emitSyntheticKeepalive = (): boolean => {
		if (finished || !keepaliveController) return false;
		lastSyntheticKeepaliveAt = Date.now();
		return tryEnqueueChunk(keepaliveController, encoder.encode(imageGenerationKeepaliveEvent()), finish);
	};
	const scheduleKeepalive = () => {
		clearKeepalive();
		if (!keepaliveController || !imageKeepaliveEnabled() || finished || policy?.keepaliveIntervalMs === undefined || policy.keepaliveIntervalMs <= 0) return;
		const delayMs = Math.max(0, policy.keepaliveIntervalMs - (Date.now() - lastSyntheticKeepaliveAt));
		keepaliveTimer = setTimeout(() => {
			if (emitSyntheticKeepalive()) scheduleKeepalive();
		}, delayMs);
	};
	const bindPolicy = (nextPolicy: RequestObservationPolicy) => {
		if (finished || policy) return;
		policy = nextPolicy;
		if (replayableLiveEvents.length > 0) {
			for (const event of replayableLiveEvents) {
				sawLiveWebSearchEvent = observeEventWithPolicy(nextPolicy, event, true, sawLiveWebSearchEvent).sawLiveWebSearchEvent;
			}
			replayableLiveEvents = [];
		}
		scheduleKeepalive();
	};
	if (isPolicyBinding(policySource)) policySource.onPolicy(bindPolicy);

	return new ReadableStream<Uint8Array>({
		start(controller) {
			keepaliveController = controller;
			scheduleKeepalive();
		},
		async pull(controller) {
			if (finished) return;
			try {
				for (;;) {
					const { value, done } = await reader.read();
					if (done) {
						finish();
						if (buffer.length > 0 && !tryEnqueueChunk(controller, encoder.encode(buffer), finish)) return;
						controller.close();
						return;
					}
					buffer += decoder.decode(value, { stream: true });
					let emitted = false;
					for (;;) {
						const delimiter = findEventDelimiter(buffer);
						if (!delimiter) break;
						const rawEvent = buffer.slice(0, delimiter.index);
						buffer = buffer.slice(delimiter.index + delimiter.length);
						const event = parseSseEvent(rawEvent);
						const observeResult = observeEventWithPolicy(policy, event, true, sawLiveWebSearchEvent);
						sawLiveWebSearchEvent = observeResult.sawLiveWebSearchEvent;
						if (!policy && observeResult.replayableEvent) replayableLiveEvents.push(observeResult.replayableEvent);
						if (!tryEnqueueChunk(controller, encoder.encode(`${rawEvent}\n\n`), finish)) return;
						emitted = true;
						if (policy?.interruptOnImageResult && isImageGenerationResultDoneEvent(rawEvent)) {
							clearTracker(policy.liveTracker);
							finish();
							if (!tryEnqueueChunk(controller, encoder.encode(INTERRUPT_DONE_EVENT), finish)) return;
							await reader.cancel().catch(() => undefined);
							controller.close();
							return;
						}
					}
					if (emitted) return;
				}
			} catch (error) {
				finish();
				failAndClearTracker(error);
				throw error;
			}
		},
		async cancel(reason) {
			finish();
			if (policy) clearTracker(policy.liveTracker);
			await reader.cancel(reason).catch(() => undefined);
		},
	});
}

export function wrapOpenAIResponsesEventIterable<T>(source: AsyncIterable<T>, policy: RequestObservationPolicy, controller?: AbortController): AsyncIterable<T | JsonRecord> {
	return {
		[Symbol.asyncIterator](): AsyncIterator<T | JsonRecord> {
			const iterator = source[Symbol.asyncIterator]();
			let finished = false;
			let upstreamNext: Promise<IteratorResult<T>> | undefined;
			let keepaliveTimer: ReturnType<typeof setTimeout> | undefined;
			let lastSyntheticKeepaliveAt = Date.now();
			let sawLiveWebSearchEvent = false;
			const imageKeepaliveEnabled = policy.enabledTools === undefined || policy.enabledTools.includes("image_generation");

			const abort = () => {
				controller?.abort();
			};
			const finish = () => {
				finished = true;
				if (keepaliveTimer) clearTimeout(keepaliveTimer);
				keepaliveTimer = undefined;
			};
			const finishAndClear = () => {
				finish();
				clearTracker(policy.liveTracker);
			};
			const failAndClear = (error: unknown) => {
				failTracker(policy.liveTracker, error);
				clearTracker(policy.liveTracker);
			};

			const nextUpstream = (): Promise<IteratorResult<T>> => {
				upstreamNext ??= iterator.next();
				return upstreamNext;
			};

			const nextKeepalive = () => new Promise<IteratorResult<JsonRecord>>(resolve => {
				if (!imageKeepaliveEnabled || policy.keepaliveIntervalMs === undefined || policy.keepaliveIntervalMs <= 0) return;
				const delayMs = Math.max(0, policy.keepaliveIntervalMs - (Date.now() - lastSyntheticKeepaliveAt));
				keepaliveTimer = setTimeout(() => {
					keepaliveTimer = undefined;
					lastSyntheticKeepaliveAt = Date.now();
					if (!finished) resolve({ value: imageGenerationKeepaliveObject(), done: false });
				}, delayMs);
			});

			return {
				async next(): Promise<IteratorResult<T | JsonRecord>> {
					if (finished) return { value: undefined, done: true };
					const upstreamResult = nextUpstream().then(
						result => ({ source: "upstream" as const, result }),
						error => ({ source: "upstreamError" as const, error }),
					);
					const keepaliveResult = nextKeepalive().then(result => ({ source: "keepalive" as const, result }));
					const outcome = await Promise.race([upstreamResult, keepaliveResult]);
					if (keepaliveTimer) clearTimeout(keepaliveTimer);
					keepaliveTimer = undefined;
					if (outcome.source === "upstreamError") {
						upstreamNext = undefined;
						finish();
						failAndClear(outcome.error);
						throw outcome.error;
					}
					if (outcome.source === "upstream") upstreamNext = undefined;
					const result = outcome.result;
					if (result.done) {
						finish();
						if (policy.observeLiveEventsInIterable !== false && !sawLiveWebSearchEvent) clearTracker(policy.liveTracker);
						return result;
					}
						if (outcome.source === "upstream" && policy.observeLiveEventsInIterable !== false) sawLiveWebSearchEvent = observeEvent(policy, result.value, true, sawLiveWebSearchEvent);
						if (policy.interruptOnImageResult && isImageGenerationResultDoneObject(result.value)) {
							finishAndClear();
							abort();
							void iterator.return?.();
						}
					return result;
				},
				async return(value?: unknown): Promise<IteratorResult<T | JsonRecord>> {
					finishAndClear();
					abort();
					void iterator.return?.();
					return { value: value as T | JsonRecord, done: true };
				},
				async throw(error?: unknown): Promise<IteratorResult<T | JsonRecord>> {
					finishAndClear();
					abort();
					void iterator.return?.();
					throw error;
				},
			};
		},
	};
}

type ObserveEventResult = {
	sawLiveWebSearchEvent: boolean;
	replayableEvent?: unknown;
};

function observeEvent(policy: RequestObservationPolicy, event: unknown, shouldCallOnEvent: boolean, sawLiveWebSearchEvent = false): boolean {
	return observeEventWithPolicy(policy, event, shouldCallOnEvent, sawLiveWebSearchEvent).sawLiveWebSearchEvent;
}

function observeEventWithPolicy(policy: RequestObservationPolicy | undefined, event: unknown, shouldCallOnEvent: boolean, sawLiveWebSearchEvent = false): ObserveEventResult {
	if (!isRecord(event)) return { sawLiveWebSearchEvent };
	const observesLiveEvent = shouldObserveLiveEvent(event);
	const shouldForward = observesLiveEvent && (event.type !== "response.completed" || sawLiveWebSearchEvent);
	if (policy && shouldCallOnEvent && shouldForward) {
		callTrackerOnEvent(policy.liveTracker, event);
	}
	const type = event.type;
	if (type === "response.completed") {
		return { sawLiveWebSearchEvent, replayableEvent: shouldForward ? event : undefined };
	} else if (type === "response.failed" || type === "error") {
		if (policy) {
			failTracker(policy.liveTracker, event.error ?? event);
			clearTracker(policy.liveTracker);
		}
	}
	return {
		sawLiveWebSearchEvent: sawLiveWebSearchEvent || isWebSearchLifecycleEvent(event),
		replayableEvent: shouldForward ? event : undefined,
	};
}

function shouldObserveLiveEvent(event: JsonRecord): boolean {
	const type = event.type;
	if (type === "response.completed" || type === "response.failed" || type === "error") return true;
	if (isWebSearchLifecycleEvent(event)) return true;
	if (type !== "response.output_item.added" && type !== "response.output_item.done") return false;
	const item = event.item;
	return isRecord(item) && item.type === "web_search_call";
}

function isWebSearchLifecycleEvent(event: JsonRecord): boolean {
	const type = event.type;
	if (typeof type === "string" && WEB_SEARCH_LIFECYCLE_EVENTS.has(type)) return true;
	if (type !== "response.output_item.added" && type !== "response.output_item.done") return false;
	const item = event.item;
	return isRecord(item) && item.type === "web_search_call";
}

function callTrackerOnEvent(tracker: ProviderToolLiveTracker | undefined, event: unknown): void {
	try {
		tracker?.onEvent(event);
	} catch {
		// Live status is best-effort and must not affect provider streams.
	}
}

function failTracker(tracker: ProviderToolLiveTracker | undefined, error: unknown): void {
	try {
		tracker?.fail(error);
	} catch {
		// Live status is best-effort and must not affect provider streams.
	}
}

function clearTracker(tracker: ProviderToolLiveTracker | undefined): void {
	try {
		tracker?.clear();
	} catch {
		// Live status is best-effort and must not affect provider streams.
	}
}

function imageGenerationKeepaliveObject(): JsonRecord {
	return {
		type: KEEPALIVE_EVENT_TYPE,
		item_id: KEEPALIVE_ITEM_ID,
		delta: "",
	};
}

function imageGenerationKeepaliveEvent(): string {
	return `data: ${JSON.stringify(imageGenerationKeepaliveObject())}\n\n`;
}

function findEventDelimiter(value: string): { index: number; length: number } | undefined {
	const lf = value.indexOf("\n\n");
	const crlf = value.indexOf("\r\n\r\n");
	if (lf === -1 && crlf === -1) return undefined;
	if (lf === -1) return { index: crlf, length: 4 };
	if (crlf === -1) return { index: lf, length: 2 };
	return crlf < lf ? { index: crlf, length: 4 } : { index: lf, length: 2 };
}

function parseSseEvent(rawEvent: string): unknown | undefined {
	const data = rawEvent
		.split(/\r?\n/)
		.filter(line => line.startsWith("data:"))
		.map(line => line.slice("data:".length).trimStart())
		.join("\n");
	if (!data || data === "[DONE]") return undefined;
	return parseJson(data);
}

function parseJson(value: string): unknown | undefined {
	try {
		return JSON.parse(value);
	} catch {
		return undefined;
	}
}

function isImageGenerationResultDoneEvent(rawEvent: string): boolean {
	const event = parseSseEvent(rawEvent);
	return isImageGenerationResultDoneObject(event);
}

function isImageGenerationResultDoneObject(event: unknown): boolean {
	if (!isRecord(event) || event.type !== "response.output_item.done") return false;
	const item = event.item;
	return isRecord(item) && item.type === "image_generation_call" && typeof item.result === "string" && item.result.length > 0;
}