#!/usr/bin/env bun

import {
	decodeWatchdogCommand,
	encodeWatchdogMessage,
	killWorkerOnSupervisorLoss,
	WATCHDOG_PROTOCOL_VERSION,
	type WatchdogStateMessage,
} from './lib/watchdog'

const MAX_IPC_BUFFER = 64 * 1024
let inputBuffer = ''
const outputClosed = false
let expectedRelease = false
let helloReceived = false
let workerPid: number | undefined
let supervisorPid: number | undefined

function parsedPid(flag: string): number | undefined {
	const value = process.argv.find((argument) => argument.startsWith(`${flag}=`))?.slice(flag.length + 1)
	if (!value || !/^[1-9][0-9]*$/.test(value)) return undefined
	const pid = Number(value)
	return Number.isSafeInteger(pid) && pid > 0 && pid !== process.pid ? pid : undefined
}

function parsedEnvironmentPid(name: string): number | undefined {
	const value = process.env[name]
	if (!value || !/^[1-9][0-9]*$/.test(value)) return undefined
	const pid = Number(value)
	return Number.isSafeInteger(pid) && pid > 0 && pid !== process.pid ? pid : undefined
}

function killWorker(): void {
	if (workerPid !== undefined) killWorkerOnSupervisorLoss(workerPid)
}

function writeState(state: WatchdogStateMessage['state']): void {
	if (outputClosed) return
	const accepted = process.stdout.write(
		encodeWatchdogMessage({
			version: WATCHDOG_PROTOCOL_VERSION,
			type: 'watchdog-state',
			state,
			pid: process.pid,
		}),
	)
	if (!accepted) throw new Error('Watchdog state pipe backpressure')
}

function fail(): void {
	if (expectedRelease) return
	killWorker()
	process.exitCode = 1
	process.stdin.pause()
}

function consumeInput(chunk: string): void {
	inputBuffer += chunk
	if (inputBuffer.length > MAX_IPC_BUFFER) {
		fail()
		return
	}
	while (true) {
		const newline = inputBuffer.indexOf('\n')
		if (newline < 0) return
		const line = inputBuffer.slice(0, newline).trim()
		inputBuffer = inputBuffer.slice(newline + 1)
		try {
			const command = decodeWatchdogCommand(line)
			if (command.type === 'hello') {
				if (supervisorPid === undefined || command.supervisorPid !== supervisorPid)
					throw new Error('Supervisor PID changed')
				helloReceived = true
				writeState('ready')
			} else {
				if (!helloReceived) throw new Error('Watchdog release arrived before handshake')
				writeState('released')
				expectedRelease = true
				process.exitCode = 0
				process.stdin.pause()
			}
		} catch {
			fail()
			return
		}
	}
}

function main(): void {
	workerPid = parsedPid('--worker-pid')
	supervisorPid = parsedPid('--supervisor-pid') ?? parsedEnvironmentPid('WATCHDOG_SUPERVISOR_PID')
	if (workerPid === undefined || supervisorPid === undefined) {
		fail()
		return
	}
	process.stdin.setEncoding('utf8')
	process.stdin.on('data', consumeInput)
	process.stdin.on('end', () => {
		if (!expectedRelease) fail()
	})
	process.stdin.on('error', fail)
	process.stdout.on('error', fail)
	process.stdout.on('close', fail)
	process.on('SIGTERM', fail)
	process.on('SIGINT', fail)
}

main()