#!/usr/bin/env bun
/**
 * vash-tv: streams a Vash broadcast channel with no PC.
 *
 * A virtual screen shows the channel's overlay in Chromium. One ffmpeg encodes
 * it; its output is copied to one small ffmpeg per destination (Twitch,
 * YouTube, …), so a destination that fails reconnects on its own without
 * touching the others. The same stream is cut into 2-second HLS segments and
 * uploaded to storage: that is the live stream on the Vash site and the
 * recording afterwards.
 *
 * It follows the channel's "Go live" switch on the site:
 *   VASH_SERVER=https://esports.vash.software VASH_CHANNEL=<id> VASH_TOKEN=<agent token>
 * Or, without Vash managing it, one overlay to one destination:
 *   OVERLAY_URL=https://…/overlay RTMP_URL=rtmp://…/<key>
 * Or relay OBS: INGEST_KEY=<secret> makes it listen on rtmp://<host>:1935/live
 * (INGEST_PORT to change) and record and restream what OBS sends, unchanged.
 * Optional: MUSIC_URL (looped audio), OVERLAY_ORIGIN (load the overlay from another
 * address, e.g. http://127.0.0.1:3102 next to Vash), DURATION (seconds, for tests).
 */
import { mkdir, readdir, readFile, rm } from 'node:fs/promises';
import type { Subprocess } from 'bun';

const env = process.env;
const server = env.VASH_SERVER?.replace(/\/$/, '');
const channelId = env.VASH_CHANNEL;
const token = env.VASH_TOKEN;
const managed = !!(server && channelId && token);
const HLS = '/tmp/vash-hls';
/** A second virtual display with osu!tourney's streamable area at its top left. */
const gameplayDisplay = env.GAMEPLAY_DISPLAY;
const gameplayWidth = Number(env.GAMEPLAY_WIDTH ?? 1280);
const gameplayHeight = Number(env.GAMEPLAY_HEIGHT ?? 720);
const say = (...parts: unknown[]) => console.log(new Date().toISOString().slice(11, 19), ...parts);

interface Destination {
	id: string;
	platform: string;
	label: string;
	url: string;
}

interface Config {
	live: boolean;
	width: number;
	height: number;
	fps: number;
	bitrateKbps: number;
	overlayUrl: string;
	record: boolean;
	destinations: Destination[];
}

if (!managed && !((env.OVERLAY_URL || env.INGEST_KEY) && env.RTMP_URL)) {
	console.error('Set VASH_SERVER, VASH_CHANNEL and VASH_TOKEN, or OVERLAY_URL and RTMP_URL.');
	process.exit(1);
}

// ── Vash API ────────────────────────────────────────────────────────────

async function api<T>(path: string, init: RequestInit = {}): Promise<T> {
	const response = await fetch(`${server}/api/broadcast/${channelId}${path}`, {
		...init,
		headers: {
			authorization: `Bearer ${token}`,
			'content-type': 'application/json',
			...init.headers
		},
		signal: AbortSignal.timeout(15_000)
	});
	if (!response.ok)
		throw new Error(
			`${path}: ${response.status} ${await response.text().catch(() => '')}`.slice(0, 300)
		);
	return response.json() as Promise<T>;
}

async function loadConfig(): Promise<Config> {
	if (!managed) {
		return {
			live: true,
			width: Number(env.WIDTH ?? 1280),
			height: Number(env.HEIGHT ?? 720),
			fps: Number(env.FPS ?? 30),
			bitrateKbps: Number((env.VIDEO_BITRATE ?? '3000k').replace(/k$/, '')),
			overlayUrl: env.OVERLAY_URL ?? '',
			record: false,
			destinations: [{ id: 'rtmp', platform: 'custom', label: 'RTMP', url: env.RTMP_URL! }]
		};
	}
	const config = await api<Config>('/stream');
	// A container next to Vash (or a dev site behind a login) can load the overlay internally.
	if (env.OVERLAY_ORIGIN)
		config.overlayUrl = config.overlayUrl.replace(
			/^https?:\/\/[^/]+/,
			env.OVERLAY_ORIGIN.replace(/\/$/, '')
		);
	return config;
}

// ── Screen ──────────────────────────────────────────────────────────────

let screen: {
	width: number;
	height: number;
	url: string;
	xvfb: Subprocess;
	browser: Subprocess | null;
} | null = null;

function spawnBrowser(url: string, width: number, height: number) {
	const background = gameplayDisplay ? 'key' : 'solid';
	const page = url.includes('?') ? `${url}&bg=${background}` : `${url}?bg=${background}`;
	return Bun.spawn(
		[
			'chromium',
			'--no-sandbox',
			'--kiosk',
			'--incognito',
			'--no-first-run',
			'--noerrdialogs',
			'--disable-infobars',
			'--disable-dev-shm-usage',
			'--disable-gpu',
			'--hide-scrollbars',
			'--force-device-scale-factor=1',
			'--autoplay-policy=no-user-gesture-required',
			'--window-position=0,0',
			`--window-size=${width},${height}`,
			page
		],
		{ env: { ...env, DISPLAY: ':99' }, stdout: 'ignore', stderr: 'ignore' }
	);
}

async function showScreen(config: Config) {
	if (
		screen &&
		screen.width === config.width &&
		screen.height === config.height &&
		screen.url === config.overlayUrl
	)
		return;
	stopScreen();
	const xvfb = Bun.spawn(
		['Xvfb', ':99', '-screen', '0', `${config.width}x${config.height}x24`, '-nolisten', 'tcp'],
		{ stdout: 'ignore', stderr: 'ignore' }
	);
	await Bun.sleep(800);
	screen = {
		width: config.width,
		height: config.height,
		url: config.overlayUrl,
		xvfb,
		browser: null
	};
	const keepBrowser = async () => {
		const current = screen;
		if (!current || current.xvfb !== xvfb) return;
		current.browser = spawnBrowser(config.overlayUrl, config.width, config.height);
		await current.browser.exited;
		// Chromium should never exit; bring it back if it does.
		if (screen === current) setTimeout(keepBrowser, 2000);
	};
	void keepBrowser();
	await Bun.sleep(5000);
	say(`Screen ${config.width}x${config.height} showing ${config.overlayUrl}`);
}

function stopScreen() {
	if (!screen) return;
	const old = screen;
	screen = null;
	old.browser?.kill();
	old.xvfb.kill();
}

// ── Encoder and outputs ────────────────────────────────────────────────

/** One consumer of the encoded stream: a destination, or the recorder. */
class Output {
	state: 'starting' | 'live' | 'reconnecting' | 'error' = 'starting';
	error: string | null = null;
	private process: Subprocess<'pipe', 'ignore', 'pipe'> | null = null;
	private since = 0;
	private queue: Uint8Array[] = [];
	private queued = 0;
	private writing = false;
	private attempts = 0;
	private stopped = false;

	constructor(
		readonly id: string,
		readonly label: string,
		private args: string[]
	) {
		this.spawn();
	}

	private spawn() {
		if (this.stopped) return;
		const process = Bun.spawn(
			[
				'ffmpeg',
				'-hide_banner',
				'-loglevel',
				'error',
				'-f',
				'mpegts',
				'-i',
				'pipe:0',
				...this.args
			],
			{
				stdin: 'pipe',
				stdout: 'ignore',
				stderr: 'pipe'
			}
		);
		this.process = process;
		this.since = Date.now();
		this.queue = [];
		this.queued = 0;
		void (async () => {
			let last = '';
			for await (const chunk of process.stderr)
				last = new TextDecoder().decode(chunk).trim().split('\n').pop() || last;
			const code = await process.exited;
			if (this.stopped || this.process !== process) return;
			this.error = last || `ffmpeg exited with ${code}`;
			this.state = this.attempts > 3 ? 'error' : 'reconnecting';
			const wait = Math.min(30_000, 2000 * 2 ** this.attempts++);
			say(`${this.label}: ${this.error}; retrying in ${Math.round(wait / 1000)}s`);
			setTimeout(() => this.spawn(), wait);
		})();
	}

	/**
	 * Queues bytes in order. Dropping any would corrupt the MPEG-TS timestamps,
	 * so a destination that falls far behind is restarted instead.
	 */
	write(chunk: Uint8Array) {
		const process = this.process;
		if (!process || this.stopped) return;
		if (this.state !== 'live' && Date.now() - this.since > 5000 && process.exitCode == null) {
			this.state = 'live';
			this.error = null;
			this.attempts = 0;
		}
		this.queue.push(chunk);
		this.queued += chunk.byteLength;
		if (this.queued > 32 * 1024 * 1024) {
			this.error = 'Destination too slow; restarting';
			this.queue = [];
			this.queued = 0;
			process.kill();
			return;
		}
		void this.drain(process);
	}

	private async drain(process: Subprocess<'pipe', 'ignore', 'pipe'>) {
		if (this.writing) return;
		this.writing = true;
		try {
			while (this.queue.length && this.process === process) {
				const chunk = this.queue.shift()!;
				this.queued -= chunk.byteLength;
				process.stdin.write(chunk);
				const flushed = process.stdin.flush();
				if (flushed instanceof Promise) await flushed;
			}
		} catch {
			// The process is exiting; its watcher restarts it.
		} finally {
			this.writing = false;
		}
	}

	stop() {
		this.stopped = true;
		try {
			this.process?.stdin.end();
		} catch {
			/* already closed */
		}
		const process = this.process;
		setTimeout(() => process?.kill(), 3000);
	}
}

let encoder: Subprocess<'ignore', 'pipe', 'pipe'> | null = null;
const outputs = new Map<string, Output>();

/** OBS (or anything) streams here instead of the built-in screen; nothing is re-encoded. */
const ingestKey = env.INGEST_KEY;
const ingestPort = Number(env.INGEST_PORT ?? 1935);

function startIngest() {
	const process = Bun.spawn(
		[
			'ffmpeg',
			'-hide_banner',
			'-loglevel',
			'error',
			'-listen',
			'1',
			'-i',
			`rtmp://0.0.0.0:${ingestPort}/live/${ingestKey}`,
			'-c',
			'copy',
			'-f',
			'mpegts',
			'pipe:1'
		],
		{ stdout: 'pipe', stderr: 'pipe' }
	);
	encoder = process;
	say(`Waiting for OBS on rtmp://<this host>:${ingestPort}/live with the ingest key`);
	void (async () => {
		for await (const chunk of process.stdout)
			for (const output of outputs.values()) output.write(chunk);
	})();
	void (async () => {
		await process.exited;
		// OBS stopped or dropped; listen again.
		if (encoder === process) encoder = null;
	})();
}

function startEncoder(config: Config) {
	if (ingestKey) return startIngest();
	const audio = env.PULSE_SOURCE
		? ['-f', 'pulse', '-i', env.PULSE_SOURCE]
		: env.MUSIC_URL
			? ['-stream_loop', '-1', '-i', env.MUSIC_URL]
			: ['-f', 'lavfi', '-i', 'anullsrc=channel_layout=stereo:sample_rate=44100'];
	const kbps = config.bitrateKbps;
	const gameplay = gameplayDisplay
		? [
				'-f',
				'x11grab',
				'-draw_mouse',
				'0',
				'-framerate',
				String(config.fps),
				'-video_size',
				`${gameplayWidth}x${gameplayHeight}`,
				'-i',
				gameplayDisplay
			]
		: [];
	const video = gameplayDisplay
		? [
				'-filter_complex',
				`[1:v]scale=${config.width}:${config.height}:force_original_aspect_ratio=decrease,pad=${config.width}:${config.height}:(ow-iw)/2:(oh-ih)/2[game];[0:v]chromakey=0x00ff00:0.04:0.0[overlay];[game][overlay]overlay=0:0:shortest=1[v]`,
				'-map',
				'[v]'
			]
		: ['-map', '0:v'];
	const process = Bun.spawn(
		[
			'ffmpeg',
			'-hide_banner',
			'-loglevel',
			'error',
			'-f',
			'x11grab',
			'-draw_mouse',
			'0',
			'-framerate',
			String(config.fps),
			'-video_size',
			`${config.width}x${config.height}`,
			'-i',
			':99',
			...gameplay,
			...audio,
			...video,
			'-map',
			gameplayDisplay ? '2:a' : '1:a',
			'-c:v',
			'libx264',
			'-preset',
			env.PRESET ?? 'veryfast',
			'-b:v',
			`${kbps}k`,
			'-maxrate',
			`${kbps}k`,
			'-bufsize',
			`${kbps * 2}k`,
			'-pix_fmt',
			'yuv420p',
			'-g',
			String(config.fps * 2),
			'-keyint_min',
			String(config.fps * 2),
			'-sc_threshold',
			'0',
			'-c:a',
			'aac',
			'-b:a',
			'160k',
			'-ar',
			'44100',
			'-f',
			'mpegts',
			'pipe:1'
		],
		{ stdout: 'pipe', stderr: 'pipe', env: { ...env, DISPLAY: ':99' } }
	);
	encoder = process;
	void (async () => {
		for await (const chunk of process.stdout)
			for (const output of outputs.values()) output.write(chunk);
	})();
	void (async () => {
		let last = '';
		for await (const chunk of process.stderr) last = new TextDecoder().decode(chunk).trim() || last;
		const code = await process.exited;
		if (encoder === process) {
			say(`Encoder stopped (${code}): ${last}`);
			encoder = null;
		}
	})();
}

function syncDestinations(config: Config) {
	const wanted = new Map(config.destinations.map((destination) => [destination.id, destination]));
	for (const [id, output] of outputs) {
		if (id === 'recording') continue;
		const destination = wanted.get(id);
		if (!destination || (output as unknown as { url?: string }).url !== destination.url) {
			output.stop();
			outputs.delete(id);
			if (destination) say(`${output.label}: settings changed, reconnecting`);
			else say(`${output.label}: removed`);
		}
	}
	for (const destination of config.destinations) {
		if (outputs.has(destination.id)) continue;
		const output = new Output(destination.id, destination.label, [
			'-c',
			'copy',
			'-f',
			'flv',
			destination.url
		]);
		(output as unknown as { url: string }).url = destination.url;
		outputs.set(destination.id, output);
		say(`${destination.label}: connecting`);
	}
}

// ── Recording (HLS to storage) ─────────────────────────────────────────

let session: { id: string; storagePrefix: string | null; startedAt: number } | null = null;
const vod: { name: string; duration: number }[] = [];
const uploaded = new Set<string>();
let lastVodUpload = 0;
let posterDone = false;

async function upload(
	files: Record<string, { body: Uint8Array | string; type: string; cache?: string }>
) {
	if (!session?.storagePrefix) return;
	const names = Object.keys(files);
	const { urls } = await api<{ urls: Record<string, string> }>('/uploads', {
		method: 'POST',
		body: JSON.stringify({ sessionId: session.id, files: names })
	});
	await Promise.all(
		names.map(async (name) => {
			const file = files[name];
			const response = await fetch(urls[name], {
				method: 'PUT',
				body: file.body,
				headers: {
					'content-type': file.type,
					...(file.cache ? { 'cache-control': file.cache } : {})
				}
			});
			if (!response.ok) throw new Error(`upload ${name}: ${response.status}`);
		})
	);
}

function vodPlaylist(ended: boolean) {
	const target = Math.ceil(Math.max(2, ...vod.map((each) => each.duration)));
	return [
		'#EXTM3U',
		'#EXT-X-VERSION:3',
		`#EXT-X-TARGETDURATION:${target}`,
		'#EXT-X-MEDIA-SEQUENCE:0',
		`#EXT-X-PLAYLIST-TYPE:${ended ? 'VOD' : 'EVENT'}`,
		...vod.flatMap((each) => [`#EXTINF:${each.duration.toFixed(3)},`, `live/${each.name}`]),
		...(ended ? ['#EXT-X-ENDLIST'] : [])
	].join('\n');
}

/** Uploads new segments, then the live playlist, and now and then the growing VOD playlist. */
async function pumpRecording(final = false) {
	if (!session?.storagePrefix) return;
	let playlist: string;
	try {
		playlist = await readFile(`${HLS}/index.m3u8`, 'utf8');
	} catch {
		return;
	}
	const entries: { name: string; duration: number }[] = [];
	const lines = playlist.split('\n');
	for (let index = 0; index < lines.length; index++) {
		const match = /^#EXTINF:([\d.]+)/.exec(lines[index]);
		if (match && lines[index + 1])
			entries.push({ name: lines[index + 1].trim(), duration: Number(match[1]) });
	}
	const fresh = entries.filter((entry) => !uploaded.has(entry.name));
	if (fresh.length) {
		const files: Record<string, { body: Uint8Array; type: string }> = {};
		for (const entry of fresh)
			files[`live/${entry.name}`] = {
				body: new Uint8Array(await readFile(`${HLS}/${entry.name}`)),
				type: 'video/mp2t'
			};
		await upload(files);
		for (const entry of fresh) {
			uploaded.add(entry.name);
			vod.push(entry);
		}
		await upload({
			'live/index.m3u8': {
				body: playlist,
				type: 'application/vnd.apple.mpegurl',
				cache: 'no-cache, max-age=1'
			}
		});
	}
	if (final || Date.now() - lastVodUpload > 30_000) {
		lastVodUpload = Date.now();
		await upload({
			'vod.m3u8': {
				body: vodPlaylist(final),
				type: 'application/vnd.apple.mpegurl',
				cache: final ? 'max-age=3600' : 'no-cache'
			}
		});
	}
	// Local segments that left the live playlist are safely stored; free the disk.
	const live = new Set(entries.map((entry) => entry.name));
	for (const name of await readdir(HLS))
		if (name.endsWith('.ts') && uploaded.has(name) && !live.has(name))
			await rm(`${HLS}/${name}`, { force: true });
	if (!posterDone && session && Date.now() - session.startedAt > 15_000 && entries.length) {
		posterDone = true;
		const latest = entries[entries.length - 1].name;
		const grab = Bun.spawn([
			'ffmpeg',
			'-hide_banner',
			'-loglevel',
			'error',
			'-y',
			'-i',
			`${HLS}/${latest}`,
			'-frames:v',
			'1',
			'-vf',
			'scale=640:-2',
			`${HLS}/poster.jpg`
		]);
		if ((await grab.exited) === 0)
			await upload({
				'poster.jpg': {
					body: new Uint8Array(await readFile(`${HLS}/poster.jpg`)),
					type: 'image/jpeg'
				}
			}).catch(() => {});
	}
}

// ── On and off air ──────────────────────────────────────────────────────

let onAir = false;

async function goLive(config: Config) {
	if (!ingestKey) await showScreen(config);
	if (managed) {
		const started = await api<{ id: string; storagePrefix: string | null }>('/stream', {
			method: 'POST',
			body: JSON.stringify({ event: 'start' })
		});
		session = { ...started, startedAt: Date.now() };
		vod.length = 0;
		uploaded.clear();
		posterDone = false;
	}
	if (config.record && session?.storagePrefix) {
		await rm(HLS, { recursive: true, force: true });
		await mkdir(HLS, { recursive: true });
		outputs.set(
			'recording',
			new Output('recording', 'Recording', [
				'-c',
				'copy',
				'-f',
				'hls',
				'-hls_time',
				'2',
				'-hls_list_size',
				'8',
				'-hls_segment_filename',
				`${HLS}/seg_%05d.ts`,
				`${HLS}/index.m3u8`
			])
		);
	}
	syncDestinations(config);
	startEncoder(config);
	onAir = true;
	say(`On air: ${config.destinations.map((each) => each.label).join(', ') || 'recording only'}`);
}

async function goOffAir() {
	onAir = false;
	encoder?.kill();
	encoder = null;
	for (const output of outputs.values()) output.stop();
	outputs.clear();
	await Bun.sleep(2500);
	await pumpRecording(true).catch((err) => say('Final upload failed:', (err as Error).message));
	if (managed && session) {
		await api('/stream', {
			method: 'POST',
			body: JSON.stringify({
				event: 'end',
				sessionId: session.id,
				durationSec: Math.round((Date.now() - session.startedAt) / 1000)
			})
		}).catch(() => {});
	}
	session = null;
	say('Off air');
}

async function report() {
	if (!managed || !session) return;
	await api('/stream', {
		method: 'POST',
		body: JSON.stringify({
			event: 'status',
			sessionId: session.id,
			durationSec: Math.round((Date.now() - session.startedAt) / 1000),
			destinations: [...outputs.values()]
				.filter((output) => output.id !== 'recording')
				.map((output) => ({
					id: output.id,
					label: output.label,
					state: output.state,
					error: output.error
				}))
		})
	}).catch((err) => say('Status report failed:', (err as Error).message));
}

let stopping = false;
async function shutdown() {
	if (stopping) return;
	stopping = true;
	if (onAir) await goOffAir();
	stopScreen();
	process.exit(0);
}
process.on('SIGTERM', shutdown);
process.on('SIGINT', shutdown);

say(managed ? `Following channel ${channelId} on ${server}` : `Streaming ${env.OVERLAY_URL}`);
if (env.DURATION) setTimeout(shutdown, Number(env.DURATION) * 1000);

let lastReport = 0;
let lastConfig = 0;
while (!stopping) {
	if (Date.now() - lastConfig >= 5000) {
		lastConfig = Date.now();
		try {
			const config = await loadConfig();
			if (config.live && !onAir) await goLive(config);
			else if (!config.live && onAir) await goOffAir();
			else if (onAir) {
				syncDestinations(config);
				if (!encoder) startEncoder(config);
			}
		} catch (err) {
			say('Vash unreachable:', (err as Error).message);
		}
	}
	if (onAir) {
		await pumpRecording().catch((err) => say('Recording upload failed:', (err as Error).message));
		if (Date.now() - lastReport > 10_000) {
			lastReport = Date.now();
			await report();
		}
	}
	await Bun.sleep(1000);
}