import type { ChatCompletionParams, StreamChunk } from '../types/provider/index.js'
import type { LLMProvider } from '../types/provider/interface.js'
import type { Logger } from '../utils/logger.js'
import { ProviderRequestError } from './errors.js'

/** Five minutes without one provider chunk is a stalled model call. */
export const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 5 * 60_000

const MAX_TIMER_DELAY_MS = 2_147_483_647

/** Resolve the shared run/router stream bound before either creates durable turn identity. */
export function resolveStreamIdleTimeoutMs(value: number | undefined): number {
	const resolved = value ?? DEFAULT_STREAM_IDLE_TIMEOUT_MS
	if (!Number.isInteger(resolved) || resolved < 0 || resolved > MAX_TIMER_DELAY_MS) {
		throw new RangeError(
			`streamIdleTimeoutMs must be an integer from 0 to ${MAX_TIMER_DELAY_MS}; received ${String(resolved)}`,
		)
	}
	return resolved
}

/**
 * Notices a stream that stopped producing without ending.
 *
 * One driver had a per-chunk watchdog, written inline, and it defaulted to
 * OFF — so zero of seven drivers re-armed on a stall unless a host set a
 * config key it had no reason to know about. The only other bound is each
 * driver's whole-REQUEST timeout, which a stream that opened successfully
 * and then went quiet does not trip: the request is fine, the bytes have
 * simply stopped.
 *
 * A turn in that state is not slow, it is stuck. It holds its budget, its
 * claim and its process, and nothing settles it until an operator
 * notices — which is the failure mode a kernel with checkpoints and
 * budgets exists to make impossible.
 *
 * ## Why a decorator and not a driver feature
 *
 * Written once here, in the shape `withProviderRetry` and
 * `withProviderFallback` already use, so it composes with them rather than
 * being reimplemented per driver. The failure is classified `network`,
 * which is the same classification the inline version produced and the one
 * retry and fallback already know how to act on: a stalled stream is
 * retried by the layer above, or the chain moves to the next member.
 *
 * ## What it cannot do
 *
 * It cannot un-emit. A stall AFTER the first chunk surfaces as an error
 * rather than a silent restart, exactly as `withProviderRetry` documents
 * for the same reason — the consumer has already emitted those deltas.
 */

export interface WithStreamIdleTimeoutOptions {
	/**
	 * Milliseconds without a chunk before the stream is treated as stalled.
	 * `0` or a non-finite value disables the watchdog and returns the
	 * provider unwrapped.
	 */
	readonly idleTimeoutMs: number
	readonly log?: Logger
	/** Test seam. Defaults to `setTimeout`. */
	readonly setTimeoutFn?: typeof setTimeout
	/** Test seam. Defaults to `clearTimeout`. */
	readonly clearTimeoutFn?: typeof clearTimeout
}

export function withStreamIdleTimeout(
	provider: LLMProvider,
	options: WithStreamIdleTimeoutOptions,
): LLMProvider {
	const { idleTimeoutMs } = options
	// Unwrapped rather than wrapped-and-inert. A disabled watchdog that
	// still races a promise per chunk costs the hottest path in the runtime
	// a timer and a closure for nothing.
	if (!Number.isFinite(idleTimeoutMs) || idleTimeoutMs <= 0) return provider

	const arm = options.setTimeoutFn ?? setTimeout
	const disarm = options.clearTimeoutFn ?? clearTimeout

	async function* chatStream(params: ChatCompletionParams): AsyncIterable<StreamChunk> {
		const callerSignal = params.signal
		callerSignal?.throwIfAborted()

		// Do not abort the caller's controller. The watchdog owns this transport
		// signal and fuses the caller's stop into it in the one safe direction.
		// Keeping it stable for the whole stream is important: drivers attach one
		// listener when the request opens and expect that signal to remain the
		// request's cancellation identity until the socket closes.
		const transportAbort = new AbortController()
		type TerminalCause =
			| { readonly kind: 'caller'; readonly reason: unknown }
			| { readonly kind: 'idle'; readonly error: ProviderRequestError }
		let terminalCause: TerminalCause | undefined
		let rejectCallerAbort: ((reason?: unknown) => void) | undefined
		const callerAbort = callerSignal
			? new Promise<never>((_resolve, reject) => {
					rejectCallerAbort = reject
				})
			: undefined
		const onCallerAbort = (): void => {
			const reason = callerSignal?.reason
			if (terminalCause === undefined) terminalCause = { kind: 'caller', reason }
			if (!transportAbort.signal.aborted) transportAbort.abort(reason)
			rejectCallerAbort?.(reason)
		}
		callerSignal?.addEventListener('abort', onCallerAbort, { once: true })

		let iteratorForCleanup: AsyncIterator<StreamChunk> | undefined
		try {
			const iterator = provider
				.chatStream({ ...params, signal: transportAbort.signal })
				[Symbol.asyncIterator]()
			iteratorForCleanup = iterator
			for (;;) {
				let timer: ReturnType<typeof setTimeout> | undefined
				let result: IteratorResult<StreamChunk>
				try {
					const idle = new Promise<never>((_resolve, reject) => {
						timer = arm(() => {
							if (terminalCause !== undefined) return
							const duration =
								idleTimeoutMs < 1_000 || idleTimeoutMs % 1_000 !== 0
									? `${idleTimeoutMs}ms`
									: `${idleTimeoutMs / 1_000}s`
							const error = new ProviderRequestError({
								kind: 'network',
								providerId: provider.id,
								detail: `stream idle for ${duration} — aborting so the turn lifecycle can settle it`,
							})
							// Latch BEFORE the transport is aborted. Some provider SDKs
							// synchronously reject their pending `next()` with a generic
							// AbortError. Retry/fallback correctly treat that shape as a
							// caller stop, so letting it escape would erase this network
							// cause and make the recovery policy unreachable.
							terminalCause = { kind: 'idle', error }
							transportAbort.abort(error)
							reject(error)
						}, idleTimeoutMs)
					})
					const pending = callerAbort
						? [iterator.next(), idle, callerAbort]
						: [iterator.next(), idle]
					result = await Promise.race(pending)
				} catch (err) {
					// The first control-flow cause wins even if the driver's own
					// abort listener rejected its pull with a different error object.
					// In particular, a generic AbortError caused by OUR watchdog is
					// still the network failure retry/fallback need to see.
					if (terminalCause?.kind === 'idle') throw terminalCause.error
					if (terminalCause?.kind === 'caller') throw terminalCause.reason
					throw err
				} finally {
					// In a `finally`, so a rejection from either side of the
					// race clears the timer. Clearing only on success leaks one
					// per stalled stream, and a leaked timer keeps the process
					// alive past the turn it belonged to.
					if (timer !== undefined) disarm(timer)
				}
				if (result.done) return
				yield result.value
			}
		} finally {
			callerSignal?.removeEventListener('abort', onCallerAbort)
			// Asked to clean up, NOT awaited — and the difference is the whole
			// case this decorator exists for. A generator stalled inside an
			// `await` is not suspended at a `yield`, so `return()` resumes
			// nothing and its promise never settles: awaiting it deadlocks on
			// precisely the stalled stream the watchdog just caught. (Found
			// that way: the first version hung this file's own test.)
			//
			// The call still matters — a driver holding a socket has no other
			// signal that nobody is reading — so it is made and abandoned,
			// with the rejection swallowed because there is no one left to
			// tell.
			try {
				void iteratorForCleanup?.return?.().catch(() => {})
			} catch {
				// A non-conforming iterator can throw before returning its cleanup
				// promise. The original stream outcome still owns this boundary.
			}
		}
	}

	// Explicit forwarding rather than object spread. A provider may be a class
	// whose methods live on its prototype; spreading it silently strips those
	// methods, and retry reads `retryDefaults` AFTER this wrapper is installed.
	return {
		get id() {
			return provider.id
		},
		get name() {
			return provider.name
		},
		get capabilities() {
			return provider.capabilities
		},
		get retryDefaults() {
			return provider.retryDefaults
		},
		chatStream,
		...(provider.listModels
			? { listModels: (signal?: AbortSignal) => provider.listModels?.(signal) }
			: {}),
		...(provider.probeCredential
			? {
					probeCredential: (signal?: AbortSignal) => provider.probeCredential?.(signal),
				}
			: {}),
		...(provider.healthCheck
			? { healthCheck: (model?: string) => provider.healthCheck?.(model) }
			: {}),
		...(provider.doctorCheck
			? { doctorCheck: (model?: string) => provider.doctorCheck?.(model) }
			: {}),
		...(provider.supportsHostedWebSearchFor
			? {
					supportsHostedWebSearchFor: (model: string, mode: 'live' | 'cached') =>
						provider.supportsHostedWebSearchFor?.(model, mode) ?? false,
				}
			: {}),
		...(provider.reasoningEffortLevelsFor
			? {
					reasoningEffortLevelsFor: (
						model: string,
						thinking?: Parameters<NonNullable<LLMProvider['reasoningEffortLevelsFor']>>[1],
					) => provider.reasoningEffortLevelsFor?.(model, thinking),
				}
			: {}),
		...(provider.reasoningEffortDefaultFor
			? {
					reasoningEffortDefaultFor: (
						model: string,
						thinking?: Parameters<NonNullable<LLMProvider['reasoningEffortDefaultFor']>>[1],
					) => provider.reasoningEffortDefaultFor?.(model, thinking),
				}
			: {}),
		...(provider.effortLevelsFor
			? {
					effortLevelsFor: (
						model: string,
						thinking?: Parameters<NonNullable<LLMProvider['effortLevelsFor']>>[1],
					) => provider.effortLevelsFor?.(model, thinking) ?? [],
				}
			: {}),
		...(provider.resolveContextWindow
			? {
					resolveContextWindow: (model: string, signal?: AbortSignal) =>
						provider.resolveContextWindow?.(model, signal) ?? Promise.resolve(undefined),
				}
			: {}),
	} as LLMProvider
}