import type { Logger } from '../../logger.ts';
import type { Readable, Writable } from 'node:stream';
import { PassThrough } from 'node:stream';
import { finished } from 'node:stream/promises';
import { z } from 'zod';
import { APIClient, PaymentRequiredError } from '../api.ts';
import { sandboxCreate } from './create.ts';
import { sandboxDestroy } from './destroy.ts';
import { executionGet } from './execution.ts';
import { sandboxGetStatus } from './getStatus.ts';
import { ExecutionCancelledError, writeAndDrain } from './util.ts';
import { SandboxRunOptionsSchema, type SandboxRunResult } from './types.ts';
import { getServiceUrls } from '../config.ts';

const timingLogsEnabled = false;
const EXECUTION_WAIT_DURATION = '5m';
const EXIT_CODE_FAST_WAIT_DURATION = '250ms';
const TERMINAL_EXECUTION_STATUSES = new Set(['completed', 'failed', 'timeout', 'cancelled']);

/**
 * Creates a Writable stream that captures all chunks to a buffer array
 * and optionally tees (forwards) them to one or more user-provided streams.
 *
 * @param chunks - Array to collect Buffer chunks into
 * @param userStreams - Optional user-provided Writable stream(s) to forward chunks to
 * @returns A Writable stream that captures and optionally forwards data
 */
function createTeeWritable(chunks: Buffer[], ...userStreams: (Writable | undefined)[]): Writable {
	const tee = new PassThrough();

	// Always capture chunks to the buffer
	tee.on('data', (chunk: Buffer) => {
		chunks.push(chunk);
	});

	// Pipe to all provided user streams with proper backpressure handling
	for (const userStream of userStreams) {
		if (userStream) {
			tee.pipe(userStream, { end: false });
		}
	}

	return tee;
}

export const SandboxRunParamsSchema = z.object({
	options: SandboxRunOptionsSchema.describe('sandbox run options'),
	orgId: z.string().optional().describe('organization id'),
	region: z.string().optional().describe('region id'),
	apiKey: z.string().optional().describe('api key'),
	signal: z.custom<AbortSignal>().optional().describe('abort signal'),
	stdin: z.custom<Readable>().optional().describe('stdin readable stream'),
	stdout: z.custom<Writable>().optional().describe('stdout writable stream'),
	stderr: z.custom<Writable>().optional().describe('stderr writable stream'),
	logger: z.custom<Logger>().optional().describe('logger instance'),
});
export type SandboxRunParams = z.infer<typeof SandboxRunParamsSchema>;

/**
 * Creates a sandbox, executes a command, and waits for completion.
 *
 * This is a high-level convenience function that handles the full lifecycle:
 * creating a sandbox, streaming I/O, polling for completion, and cleanup.
 *
 * @param client - The API client to use for the request
 * @param params - Parameters including command options, I/O streams, and timeout settings
 * @returns The run result including exit code and duration
 * @throws {SandboxResponseError} If sandbox creation fails, execution times out, or is cancelled
 */
export async function sandboxRun(
	client: APIClient,
	params: SandboxRunParams
): Promise<SandboxRunResult> {
	const { options, orgId, region, apiKey, signal, stdin, stdout, stderr, logger } = params;
	const started = Date.now();
	if (timingLogsEnabled) console.error(`[TIMING] +0ms: sandbox run started`);

	let stdinStreamId: string | undefined;
	let stdinStreamUrl: string | undefined;

	// Handle stdin stream configuration:
	// - If stdin is "ignore", pass it through to skip stdin handling on server
	// - If stdin is an explicit stream ID, use it directly
	// - If stdin readable is provided, create a stream for it
	const stdinConfig = options.stream?.stdin;
	if (stdinConfig === 'ignore') {
		stdinStreamId = 'ignore';
		logger?.debug('stdin explicitly ignored');
	} else if (stdinConfig && stdinConfig !== 'ignore') {
		// User provided an explicit stream ID
		stdinStreamId = stdinConfig;
		logger?.debug('using provided stdin stream ID: %s', stdinStreamId);
	} else if (stdin && region && apiKey) {
		const streamResult = await createStdinStream(region, apiKey, orgId, logger);
		stdinStreamId = streamResult.id;
		stdinStreamUrl = streamResult.url;
		logger?.debug('created stdin stream: %s', stdinStreamId);
	}

	const createResponse = await sandboxCreate(client, {
		options: {
			...options,
			command: {
				exec: options.command.exec,
				files: options.command.files,
				mode: 'oneshot',
			},
			stream: {
				...options.stream,
				stdin: stdinStreamId,
			},
		},
		orgId,
	});

	const sandboxId = createResponse.sandboxId;
	const stdoutStreamUrl = createResponse.stdoutStreamUrl;
	const stderrStreamUrl = createResponse.stderrStreamUrl;

	logger?.debug(
		'sandbox created: %s, stdoutUrl: %s, stderrUrl: %s',
		sandboxId,
		stdoutStreamUrl ?? 'none',
		stderrStreamUrl ?? 'none'
	);
	if (timingLogsEnabled)
		console.error(`[TIMING] +${Date.now() - started}ms: sandbox created (${sandboxId})`);

	const abortController = new AbortController();
	const streamPromises: Promise<void>[] = [];

	// Create capture buffers for stdout/stderr
	const stdoutChunks: Buffer[] = [];
	const stderrChunks: Buffer[] = [];

	try {
		// Start stdin streaming if we have stdin and a stream URL
		if (stdin && stdinStreamUrl && apiKey) {
			const stdinPromise = streamStdinToUrl(
				stdin,
				stdinStreamUrl,
				apiKey,
				abortController.signal,
				logger
			);
			streamPromises.push(stdinPromise);
		}

		// Check if stdout and stderr are the same stream (combined output)
		const isCombinedOutput =
			stdoutStreamUrl && stderrStreamUrl && stdoutStreamUrl === stderrStreamUrl;

		if (isCombinedOutput) {
			// Stream combined output to stdout only to avoid duplicates
			if (stdoutStreamUrl) {
				logger?.debug('using combined output stream (stdout === stderr)');
				const teeStream = createTeeWritable(stdoutChunks, stdout);
				const combinedPromise = streamUrlToWritable(
					stdoutStreamUrl,
					teeStream,
					abortController.signal,
					logger,
					started
				);
				streamPromises.push(combinedPromise);
			}
		} else {
			// Start stdout streaming with capture
			if (stdoutStreamUrl) {
				const teeStream = createTeeWritable(stdoutChunks, stdout);
				const stdoutPromise = streamUrlToWritable(
					stdoutStreamUrl,
					teeStream,
					abortController.signal,
					logger,
					started
				);
				streamPromises.push(stdoutPromise);
			}

			// Start stderr streaming with capture
			if (stderrStreamUrl) {
				const teeStream = createTeeWritable(stderrChunks, stderr);
				const stderrPromise = streamUrlToWritable(
					stderrStreamUrl,
					teeStream,
					abortController.signal,
					logger,
					started
				);
				streamPromises.push(stderrPromise);
			}
		}

		// Wait for execution completion in parallel with stream consumption. The old
		// flow waited for stream EOF first and only then started polling for the
		// final exit code, which adds avoidable tail latency now that create returns
		// an execution ID immediately for oneshot sandboxes.
		let finalExecution:
			| {
					exitCode?: number;
					status: string;
			  }
			| undefined;
		if (createResponse.executionId) {
			logger?.debug(
				'waiting for execution %s and %d stream(s) in parallel',
				createResponse.executionId,
				streamPromises.length
			);
			const completionPromise = waitForRunCompletion(
				client,
				sandboxId,
				createResponse.executionId,
				orgId,
				signal,
				logger,
				started
			);

			finalExecution = signal
				? await raceWithAbort(completionPromise, signal, abortController, sandboxId)
				: await completionPromise;
			await waitForStreamsToDrain(streamPromises, signal, abortController, sandboxId);
		} else {
			logger?.debug(
				'missing executionId on create response, falling back to stream-first completion'
			);
			await waitForStreamsToDrain(streamPromises, signal, abortController, sandboxId);
		}

		if (timingLogsEnabled)
			console.error(`[TIMING] +${Date.now() - started}ms: completion wait finished`);
		logger?.debug('completion wait finished, resolving final exit code');

		// Stream EOF means the sandbox is done — hadron only closes streams after the
		// container exits. Poll for the exit code with retries because the lifecycle
		// event (carrying the exit code) may still be in flight to Catalyst when the
		// stream completes.
		//
		// Hadron drains container logs for up to 5s after exit, then closes the
		// stream, then sends the lifecycle event in a goroutine. So the exit code
		// typically arrives at Catalyst 5–7s after the container exits. We use a
		// linear 1s polling interval (not exponential backoff) so we don't overshoot
		// the window — 15 attempts × 1s = 15s total, which comfortably covers the
		// drain + lifecycle propagation delay.
		let exitCode = finalExecution?.exitCode ?? 0;
		const statusPollStart = Date.now();
		let shouldWaitForSandboxStatus = finalExecution?.exitCode == null;
		let sandboxStatusReconciled = false;
		if (finalExecution?.exitCode == null) {
			if (createResponse.executionId && finalExecution?.status === 'completed') {
				try {
					const execution = await executionGet(client, {
						executionId: createResponse.executionId,
						orgId,
						wait: EXIT_CODE_FAST_WAIT_DURATION,
						signal,
					});
					if (execution.exitCode != null) {
						exitCode = execution.exitCode;
						finalExecution.exitCode = execution.exitCode;
						shouldWaitForSandboxStatus = false;
						logger?.debug(
							'[run] exit code %d found from fast execution retry (+%dms)',
							exitCode,
							Date.now() - statusPollStart
						);
					}
				} catch (err) {
					if (!(err instanceof DOMException && err.name === 'AbortError')) {
						logger?.debug(
							'[run] fast execution exit code retry failed (+%dms): %s',
							Date.now() - statusPollStart,
							err
						);
					}
				}
			}
		}
		if (shouldWaitForSandboxStatus) {
			try {
				const sandboxStatus = await sandboxGetStatus(client, {
					sandboxId,
					orgId,
					waitForStatus: ['terminated', 'failed'],
					waitMs: 15000,
				});
				if (sandboxStatus.exitCode != null) {
					exitCode = sandboxStatus.exitCode;
					sandboxStatusReconciled = true;
					logger?.debug(
						'[run] exit code %d found after server-side wait (+%dms)',
						exitCode,
						Date.now() - statusPollStart
					);
				} else if (sandboxStatus.status === 'failed') {
					exitCode = 1;
					sandboxStatusReconciled = true;
					logger?.debug(
						'[run] sandbox failed after server-side wait (+%dms)',
						Date.now() - statusPollStart
					);
				} else if (sandboxStatus.status === 'terminated') {
					sandboxStatusReconciled = true;
					logger?.debug(
						'[run] sandbox terminated without exit code after server-side wait (+%dms)',
						Date.now() - statusPollStart
					);
				} else {
					logger?.debug(
						'[run] sandbox status wait expired with status=%s (+%dms)',
						sandboxStatus.status,
						Date.now() - statusPollStart
					);
				}
			} catch (err) {
				if (!(err instanceof DOMException && err.name === 'AbortError')) {
					logger?.debug(
						'[run] sandboxGetStatus server-side wait failed (+%dms): %s',
						Date.now() - statusPollStart,
						err
					);
				}
			}
		}
		if (
			finalExecution &&
			finalExecution?.exitCode == null &&
			finalExecution?.status !== 'completed' &&
			!sandboxStatusReconciled
		) {
			exitCode = 1;
			logger?.debug(
				'[run] using fallback exit code 1 for terminal status=%s after sandbox status reconciliation failed',
				finalExecution?.status
			);
		}
		if (exitCode === 0) {
			if (finalExecution?.exitCode != null) {
				logger?.debug('[run] using execution exit code 0 from long-poll result');
			} else {
				logger?.debug(
					'[run] exit code wait finished with default 0 (+%dms)',
					Date.now() - statusPollStart
				);
			}
		}

		if (timingLogsEnabled)
			console.error(
				`[TIMING] +${Date.now() - started}ms: sandboxGet complete (exit: ${exitCode})`
			);

		// Build captured output strings
		const capturedStdout = Buffer.concat(stdoutChunks).toString('utf-8');
		const capturedStderr = isCombinedOutput
			? capturedStdout
			: Buffer.concat(stderrChunks).toString('utf-8');

		return {
			sandboxId,
			exitCode,
			durationMs: Date.now() - started,
			stdout: capturedStdout,
			stderr: capturedStderr,
		};
	} catch (error) {
		abortController.abort();
		try {
			await sandboxDestroy(client, { sandboxId, orgId });
		} catch {
			// Ignore cleanup errors
		}
		throw error;
	}
}

async function waitForRunCompletion(
	client: APIClient,
	sandboxId: string,
	executionId: string,
	orgId: string | undefined,
	signal: AbortSignal | undefined,
	logger: Logger | undefined,
	started: number
): Promise<{ exitCode?: number; status: string }> {
	const completionAbortController = new AbortController();
	let onAbort: (() => void) | undefined;
	if (signal) {
		onAbort = () => completionAbortController.abort(signal.reason);
		if (signal.aborted) {
			onAbort();
		} else {
			signal.addEventListener('abort', onAbort, { once: true });
		}
	}

	try {
		const completionSignal = completionAbortController.signal;
		const executionPromise = waitForExecutionCompletion(
			client,
			executionId,
			orgId,
			completionSignal,
			logger,
			started
		);
		const statusPromise = waitForSandboxStatusCompletion(
			client,
			sandboxId,
			orgId,
			completionSignal,
			logger,
			started
		).catch((err) => {
			if (completionSignal.aborted) {
				throw err;
			}
			logger?.debug('[run] sandbox status completion wait failed: %s', err);
			return new Promise<never>(() => {});
		});

		const result = await Promise.race([executionPromise, statusPromise]);
		return result;
	} finally {
		if (onAbort && signal) {
			signal.removeEventListener('abort', onAbort);
		}
	}
}

async function waitForExecutionCompletion(
	client: APIClient,
	executionId: string,
	orgId: string | undefined,
	signal: AbortSignal | undefined,
	logger: Logger | undefined,
	started: number
): Promise<{ exitCode?: number; status: string }> {
	while (true) {
		if (signal?.aborted) {
			throw new DOMException('Aborted', 'AbortError');
		}

		const result = await executionGet(client, {
			executionId,
			orgId,
			wait: EXECUTION_WAIT_DURATION,
			signal,
		});
		logger?.debug(
			'[run] execution wait: id=%s status=%s exit=%s +%dms',
			executionId,
			result.status,
			result.exitCode ?? 'undefined',
			Date.now() - started
		);

		if (TERMINAL_EXECUTION_STATUSES.has(result.status)) {
			return {
				exitCode: result.exitCode,
				status: result.status,
			};
		}
	}
}

async function waitForSandboxStatusCompletion(
	client: APIClient,
	sandboxId: string,
	orgId: string | undefined,
	signal: AbortSignal | undefined,
	logger: Logger | undefined,
	started: number
): Promise<{ exitCode?: number; status: string }> {
	while (true) {
		if (signal?.aborted) {
			throw new DOMException('Aborted', 'AbortError');
		}

		const result = await sandboxGetStatus(client, {
			sandboxId,
			orgId,
			waitForStatus: ['idle', 'terminated', 'failed'],
			waitMs: 300000,
			signal,
		});
		logger?.debug(
			'[run] sandbox status wait: sandbox=%s status=%s exit=%s +%dms',
			sandboxId,
			result.status,
			result.exitCode ?? 'undefined',
			Date.now() - started
		);

		if (result.exitCode != null) {
			return {
				exitCode: result.exitCode,
				status: 'completed',
			};
		}
		if (result.status === 'failed') {
			return {
				exitCode: 1,
				status: 'failed',
			};
		}
		if (result.status === 'terminated') {
			return {
				status: 'completed',
			};
		}

		await new Promise((resolve) => setTimeout(resolve, 25));
	}
}

async function waitForStreamsToDrain(
	streamPromises: Promise<void>[],
	signal: AbortSignal | undefined,
	abortController: AbortController,
	sandboxId: string
): Promise<void> {
	if (streamPromises.length === 0) {
		return;
	}

	if (signal) {
		let onAbort: (() => void) | undefined;
		try {
			await Promise.race([
				Promise.allSettled(streamPromises).then(() => undefined),
				new Promise<never>((_, reject) => {
					onAbort = () => {
						abortController.abort();
						reject(
							new ExecutionCancelledError({
								message: 'Sandbox execution cancelled',
								sandboxId,
							})
						);
					};
					if (signal.aborted) {
						onAbort();
					} else {
						signal.addEventListener('abort', onAbort, { once: true });
					}
				}),
			]);
		} finally {
			if (onAbort) {
				signal.removeEventListener('abort', onAbort);
			}
		}
		return;
	}

	await Promise.allSettled(streamPromises);
}

async function raceWithAbort<T>(
	promise: Promise<T>,
	signal: AbortSignal,
	abortController: AbortController,
	sandboxId: string
): Promise<T> {
	let onAbort: (() => void) | undefined;
	try {
		return await Promise.race([
			promise,
			new Promise<never>((_, reject) => {
				onAbort = () => {
					abortController.abort();
					reject(
						new ExecutionCancelledError({
							message: 'Sandbox execution cancelled',
							sandboxId,
						})
					);
				};
				if (signal.aborted) {
					onAbort();
				} else {
					signal.addEventListener('abort', onAbort, { once: true });
				}
			}),
		]);
	} finally {
		if (onAbort) {
			signal.removeEventListener('abort', onAbort);
		}
	}
}

async function createStdinStream(
	region: string,
	apiKey: string,
	orgId?: string,
	logger?: Logger
): Promise<{ id: string; url: string }> {
	const urls = getServiceUrls(region);
	const streamBaseUrl = urls.stream;

	// Build URL with orgId query param for CLI token validation
	const queryParams = new URLSearchParams();
	if (orgId) {
		queryParams.set('orgId', orgId);
	}
	const queryString = queryParams.toString();
	const url = `${streamBaseUrl}${queryString ? `?${queryString}` : ''}`;
	logger?.trace('creating stdin stream: %s', url);

	const response = await fetch(url, {
		method: 'POST',
		headers: {
			'Content-Type': 'application/json',
			Authorization: `Bearer ${apiKey}`,
		},
		body: JSON.stringify({
			name: `sandbox-stdin-${Date.now()}`,
		}),
	});

	if (!response.ok) {
		if (response.status === 402) {
			throw new PaymentRequiredError({
				url: url,
			});
		}
		throw new Error(`Failed to create stdin stream: ${response.status} ${response.statusText}`);
	}

	const data = (await response.json()) as { id: string };
	logger?.debug('created stdin stream: %s', data.id);

	// Include orgId in the URL for subsequent PUT requests (needed for CLI token auth)
	const putQueryString = orgId ? `?orgId=${encodeURIComponent(orgId)}` : '';
	return {
		id: data.id,
		url: `${streamBaseUrl}/${data.id}${putQueryString}`,
	};
}

async function streamStdinToUrl(
	stdin: Readable,
	url: string,
	apiKey: string,
	signal: AbortSignal,
	logger?: Logger
): Promise<void> {
	try {
		logger?.debug('streaming stdin to: %s', url);

		// Convert Node.js Readable to a web ReadableStream for fetch body
		let controllerClosed = false;
		const webStream = new ReadableStream({
			start(controller) {
				stdin.on('data', (chunk: Buffer) => {
					if (!signal.aborted && !controllerClosed) {
						controller.enqueue(chunk);
					}
				});
				stdin.on('end', () => {
					if (!controllerClosed) {
						controllerClosed = true;
						controller.close();
					}
				});
				stdin.on('error', (err) => {
					if (!controllerClosed) {
						controllerClosed = true;
						controller.error(err);
					}
				});
				signal.addEventListener('abort', () => {
					if (!controllerClosed) {
						controllerClosed = true;
						controller.close();
					}
				});
			},
		});

		const response = await fetch(url, {
			method: 'PUT',
			headers: {
				Authorization: `Bearer ${apiKey}`,
			},
			body: webStream,
			signal,
			duplex: 'half',
		} as RequestInit);

		if (!response.ok) {
			logger?.debug('stdin stream PUT failed: %d', response.status);
		} else {
			logger?.debug('stdin stream completed');
		}
	} catch (err) {
		if (err instanceof Error && err.name === 'AbortError') {
			logger?.debug('stdin stream aborted (expected on completion)');
		} else {
			logger?.debug('stdin stream error: %s', err);
		}
	}
}

async function streamUrlToWritable(
	url: string,
	writable: Writable,
	signal: AbortSignal,
	logger?: Logger,
	_started?: number
): Promise<void> {
	const streamStart = Date.now();
	try {
		// Signal to Pulse that this is a v2 stream so it waits for v2 metadata
		// instead of falling back to the legacy download path on a short timeout.
		const v2Url = new URL(url);
		v2Url.searchParams.set('v', '2');
		logger?.debug('[stream] fetching: %s', v2Url.href);
		const response = await fetch(v2Url.href, { signal });
		logger?.debug(
			'[stream] response status=%d in %dms',
			response.status,
			Date.now() - streamStart
		);

		if (!response.ok || !response.body) {
			logger?.debug('[stream] not ok or no body (status=%d) — returning empty', response.status);
			return;
		}

		const reader = response.body.getReader();
		let chunks = 0;
		let totalBytes = 0;

		// Read until EOF - Pulse will block until data is available
		while (true) {
			const { done, value } = await reader.read();
			if (done) {
				logger?.debug(
					'[stream] EOF after %dms (%d chunks, %d bytes)',
					Date.now() - streamStart,
					chunks,
					totalBytes
				);
				break;
			}

			if (value) {
				chunks++;
				totalBytes += value.length;
				if (chunks <= 3 || chunks % 100 === 0) {
					logger?.debug(
						'[stream] chunk #%d: %d bytes (total: %d bytes, +%dms)',
						chunks,
						value.length,
						totalBytes,
						Date.now() - streamStart
					);
				}
				await writeAndDrain(writable, value);
			}
		}
		// Signal end-of-stream to the tee/pipe chain so downstream
		// consumers (e.g. process.stdout pipe) know no more data is coming.
		writable.end();
		if ('once' in writable) {
			await finished(writable as NodeJS.WritableStream).catch(() => {
				// Ignore finish errors here; the main read/write path already
				// reported meaningful stream errors.
			});
		}
	} catch (err) {
		if (err instanceof Error && err.name === 'AbortError') {
			logger?.debug('[stream] aborted after %dms', Date.now() - streamStart);
			return;
		}
		logger?.debug('[stream] error after %dms: %s', Date.now() - streamStart, err);
	}
}
