#!/usr/bin/env node
// Standalone broker for a pi-chrome companion extension on a shared host (see FORK.md).
//
// Upstream's bridge lives inside whichever pi session started first on the machine running the
// browser: it owns 127.0.0.1:17318, the companion extension long-polls it, and later sessions on
// that same machine forward through it. That cannot serve pi sessions on other hosts - the
// extension only ever polls its own loopback, and the owner disappears with its session.
//
// This process takes the owner's place, permanently, on the browser's host:
//
//   extension listener  EXT_HOST:PORT (loopback)  /next, /result, /status, /command
//   peer listener       PEER_HOST:PORT (private)  /command, /status - nothing else
//
// Keeping /next and /result off the peer listener means nothing on the network can drain the
// extension's command queue or forge its results. /command needs the shared bearer token on
// both listeners, because loopback here is shared with everything else running as the user and
// the browser behind it is signed into real accounts.
//
// The HTTP contract towards the extension is upstream's, unchanged: the extension is loaded as
// shipped and has no idea it is talking to anything but a pi session.
//
// Node builtins only, so it runs straight from a checkout with no install step.

import { createServer } from "node:http";
import { readFileSync } from "node:fs";
import { execSync } from "node:child_process";
import { createHash, timingSafeEqual } from "node:crypto";
import { dirname, resolve } from "node:path";
import { fileURLToPath } from "node:url";
import { assertUploadsStaged } from "./upload-guard.mjs";

const PKG_PATH = resolve(dirname(fileURLToPath(import.meta.url)), "..", "package.json");

const PORT = Number(process.env.PI_CHROME_BROKER_PORT ?? "17318");
const EXT_HOST = process.env.PI_CHROME_BROKER_EXT_HOST ?? "127.0.0.1";
const PEER_HOST = process.env.PI_CHROME_BROKER_PEER_HOST?.trim() || undefined;

const DEFAULT_TIMEOUT_MS = 30_000;
const MAX_TIMEOUT_MS = 10 * 60_000;
const LONG_POLL_MS = 25_000;
// Results carry screenshots and full-page tiles as base64, so this is sized for those.
const MAX_BODY_BYTES = 64 * 1024 * 1024;

function log(level, msg, extra = {}) {
	const line = JSON.stringify({ t: new Date().toISOString(), level, msg, ...extra });
	if (level === "error") console.error(line);
	else console.log(line);
}

function readVersion() {
	try {
		return JSON.parse(readFileSync(PKG_PATH, "utf8")).version ?? "0.0.0-dev";
	} catch {
		return "0.0.0-dev";
	}
}

// PI_CHROME_TOKEN, or the output of PI_CHROME_TOKEN_COMMAND for a secret store - the same token
// every client sends. There is no default: the browser behind this is signed into real accounts.
function loadToken() {
	const fromEnv = process.env.PI_CHROME_TOKEN?.trim();
	if (fromEnv) return fromEnv;
	const command = process.env.PI_CHROME_TOKEN_COMMAND?.trim();
	if (!command) throw new Error("the broker needs a token: set PI_CHROME_TOKEN, or PI_CHROME_TOKEN_COMMAND to read one from a secret store");
	const token = execSync(command, { encoding: "utf8", stdio: ["ignore", "pipe", "pipe"] }).trim();
	if (!token) throw new Error(`PI_CHROME_TOKEN_COMMAND produced nothing: ${command}`);
	return token;
}

const TOKEN_DIGEST = createHash("sha256").update(loadToken()).digest();

// Compare digests rather than the raw strings: timingSafeEqual needs equal lengths, and the
// length of a guess must not leak either.
function hasValidToken(request) {
	const raw = String(request.headers.authorization ?? "");
	const match = /^Bearer\s+(.+)$/i.exec(raw);
	if (!match) return false;
	return timingSafeEqual(createHash("sha256").update(match[1].trim()).digest(), TOKEN_DIGEST);
}

// Same origin rules as upstream's bridge (index.ts), so the extension sees identical behaviour.
function corsHeadersFor(request) {
	const origin = String(request.headers.origin ?? "");
	if (!origin.startsWith("chrome-extension://")) return {};
	return {
		"access-control-allow-origin": origin,
		"access-control-allow-methods": "GET,POST,OPTIONS",
		"access-control-allow-headers": "content-type",
		"access-control-expose-headers": "x-pi-chrome-version",
		"vary": "origin",
	};
}

function isBrowserOriginAllowed(request) {
	const origin = String(request.headers.origin ?? "");
	if (origin) return origin.startsWith("chrome-extension://");
	const secFetchSite = String(request.headers["sec-fetch-site"] ?? "");
	return !secFetchSite || secFetchSite === "none" || secFetchSite === "same-origin";
}

function isLocalProcessRequest(request) {
	return !request.headers.origin && !request.headers["sec-fetch-site"];
}

function sendJson(response, status, body, extraHeaders = {}) {
	if (response.writableEnded) return;
	response.writeHead(status, {
		"content-type": "application/json; charset=utf-8",
		"cache-control": "no-store",
		...extraHeaders,
	});
	response.end(JSON.stringify(body));
}

function readBody(request) {
	return new Promise((resolveBody, rejectBody) => {
		const chunks = [];
		let size = 0;
		request.on("data", (chunk) => {
			size += chunk.length;
			if (size > MAX_BODY_BYTES) {
				rejectBody(new Error(`request body over ${MAX_BODY_BYTES} bytes`));
				request.destroy();
				return;
			}
			chunks.push(Buffer.from(chunk));
		});
		request.on("end", () => resolveBody(Buffer.concat(chunks).toString("utf8")));
		request.on("error", rejectBody);
	});
}

// --- command queue: a port of upstream's ChromeProfileBridge server half ----------------------

const queue = [];
const waiters = [];
const pending = new Map();
let lastSeenAt;
let clientName;
const startedAt = Date.now();

function enqueue(command) {
	const waiter = waiters.shift();
	if (waiter) waiter(command);
	else queue.push(command);
}

function dropQueued(id) {
	const index = queue.findIndex((queued) => queued.id === id);
	if (index !== -1) queue.splice(index, 1);
}

function waitForCommand(timeoutMs, registerWaiter) {
	return new Promise((resolveWait) => {
		let settled = false;
		const waiter = (command) => {
			if (settled) return;
			settled = true;
			clearTimeout(timer);
			const index = waiters.indexOf(waiter);
			if (index !== -1) waiters.splice(index, 1);
			resolveWait(command);
		};
		const timer = setTimeout(() => waiter(undefined), timeoutMs);
		waiters.push(waiter);
		registerWaiter(waiter);
	});
}

// Upstream's three-way diagnosis, so a client sees the same actionable timeout text.
function timeoutMessage(entry, timeoutMs) {
	const pollAgeMs = lastSeenAt === undefined ? undefined : Date.now() - lastSeenAt;
	if (entry?.deliveredAt) {
		return `Timed out after ${timeoutMs}ms: the browser extension received the command but never returned a result. The action may be long-running, or the result post failed. Retry; if it persists, reload 'Pi Chrome Connector' on the browser host.`;
	}
	if (pollAgeMs === undefined || pollAgeMs > 60_000) {
		return `Timed out after ${timeoutMs}ms: the browser extension is not polling the broker (last seen ${pollAgeMs === undefined ? "never" : Math.round(pollAgeMs / 1000) + "s ago"}). The browser on the broker's host is closed, or 'Pi Chrome Connector' is not loaded in it.`;
	}
	return `Timed out after ${timeoutMs}ms: the browser extension is polling (last seen ${Math.round(pollAgeMs / 1000)}s ago) but did not pick up this command in time. Retry; if it persists, reload 'Pi Chrome Connector' on the browser host.`;
}

function runCommand(action, params, timeoutMs, onAbort) {
	const id = `${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}`;
	const command = { id, action, params };
	return new Promise((resolveCommand, rejectCommand) => {
		const timer = setTimeout(() => {
			const entry = pending.get(id);
			pending.delete(id);
			dropQueued(id);
			rejectCommand(new Error(timeoutMessage(entry, timeoutMs)));
		}, timeoutMs);
		pending.set(id, { command, resolve: resolveCommand, reject: rejectCommand, timer });
		// A client that hangs up before delivery should not leave its command to run later
		// against a browser nobody is waiting on. Once delivered, it is the extension's.
		onAbort(() => {
			const entry = pending.get(id);
			if (!entry || entry.deliveredAt) return;
			clearTimeout(entry.timer);
			pending.delete(id);
			dropQueued(id);
			rejectCommand(new Error("client went away"));
		});
		enqueue(command);
	});
}

function status() {
	return {
		mode: "broker",
		version: readVersion(),
		startedAt,
		connected: lastSeenAt !== undefined && Date.now() - lastSeenAt < 5 * 60_000,
		lastSeenAt,
		clientName,
		queuedCommands: queue.length,
		pendingCommands: pending.size,
	};
}

async function handleCommand(request, response) {
	if (!isLocalProcessRequest(request)) {
		sendJson(response, 403, { ok: false, error: "Chrome commands are accepted only from local Pi processes" });
		return;
	}
	if (!hasValidToken(request)) {
		sendJson(response, 401, { ok: false, error: "pi-chrome broker rejected the bearer token (check PI_CHROME_TOKEN on both ends)" });
		return;
	}
	let body;
	try {
		body = JSON.parse(await readBody(request));
	} catch (error) {
		sendJson(response, 400, { ok: false, error: `Bad command body: ${error.message}` });
		return;
	}
	if (!body?.action) {
		sendJson(response, 400, { ok: false, error: "Missing command action" });
		return;
	}
	try {
		assertUploadsStaged(body.action, body.params, process.env.PI_CHROME_UPLOAD_DIR?.trim() || undefined);
	} catch (error) {
		log("warn", "command", { action: body.action, peer: request.socket.remoteAddress, ok: false, refused: "upload outside the staging directory" });
		sendJson(response, 403, { ok: false, error: error.message });
		return;
	}
	const timeoutMs = Math.min(Math.max(Number(body.timeoutMs) || DEFAULT_TIMEOUT_MS, 1_000), MAX_TIMEOUT_MS);
	const started = Date.now();
	const peer = request.socket.remoteAddress;
	try {
		const result = await runCommand(body.action, body.params ?? {}, timeoutMs, (abort) => {
			response.once("close", () => {
				if (!response.writableEnded) abort();
			});
		});
		// Log the action, never its params: those carry typed text, form values and page data.
		log("info", "command", { action: body.action, peer, ok: true, ms: Date.now() - started });
		sendJson(response, 200, { ok: true, result });
	} catch (error) {
		log("warn", "command", { action: body.action, peer, ok: false, ms: Date.now() - started, error: error.message });
		sendJson(response, 504, { ok: false, error: error.message });
	}
}

async function handleNext(request, response, url) {
	if (!isBrowserOriginAllowed(request)) {
		sendJson(response, 403, { ok: false, error: "browser origin not allowed" });
		return;
	}
	if (lastSeenAt === undefined || Date.now() - lastSeenAt > 60_000) {
		log("info", "extension polling", { name: url.searchParams.get("name") ?? undefined });
	}
	lastSeenAt = Date.now();
	clientName = url.searchParams.get("name") ?? undefined;
	let aborted = false;
	let activeWaiter;
	response.once("close", () => {
		if (response.writableEnded) return;
		aborted = true;
		const index = activeWaiter ? waiters.indexOf(activeWaiter) : -1;
		if (index !== -1) waiters.splice(index, 1);
	});
	let command = queue.shift();
	if (!command) {
		command = await waitForCommand(LONG_POLL_MS, (waiter) => {
			activeWaiter = waiter;
		});
	}
	if (aborted) {
		// The long-poll died before delivery: put the command back for the next poll.
		if (command) queue.unshift(command);
		return;
	}
	if (command) {
		const entry = pending.get(command.id);
		if (entry) entry.deliveredAt = Date.now();
	}
	const version = readVersion();
	sendJson(
		response,
		200,
		command ? { type: "command", command, expectedExtensionVersion: version } : { type: "none", expectedExtensionVersion: version },
		{ ...corsHeadersFor(request), "x-pi-chrome-version": version },
	);
}

async function handleResult(request, response) {
	if (!isBrowserOriginAllowed(request)) {
		sendJson(response, 403, { ok: false, error: "browser origin not allowed" });
		return;
	}
	lastSeenAt = Date.now();
	const corsHeaders = corsHeadersFor(request);
	let result;
	try {
		result = JSON.parse(await readBody(request));
	} catch (error) {
		sendJson(response, 400, { ok: false, error: `Bad result body: ${error.message}` }, corsHeaders);
		return;
	}
	const entry = pending.get(result?.id);
	if (!entry) {
		sendJson(response, 404, { ok: false, error: "unknown command id" }, corsHeaders);
		return;
	}
	clearTimeout(entry.timer);
	pending.delete(result.id);
	if (result.ok) entry.resolve(result.result);
	else entry.reject(new Error(result.error ?? "Chrome extension command failed"));
	sendJson(response, 200, { ok: true }, corsHeaders);
}

async function handle(listener, request, response) {
	const url = new URL(request.url ?? "/", "http://broker");
	const route = `${request.method} ${url.pathname}`;

	if (route === "POST /command") return handleCommand(request, response);

	if (route === "GET /status") {
		if (listener === "peer" && !hasValidToken(request)) {
			sendJson(response, 401, { ok: false, error: "bearer token required" });
			return;
		}
		sendJson(response, 200, status());
		return;
	}

	if (listener === "ext") {
		if (request.method === "OPTIONS") {
			if (!isBrowserOriginAllowed(request)) {
				sendJson(response, 403, { ok: false, error: "browser origin not allowed" });
				return;
			}
			sendJson(response, 200, { ok: true }, corsHeadersFor(request));
			return;
		}
		if (route === "GET /next") return handleNext(request, response, url);
		if (route === "POST /result") return handleResult(request, response);
	}

	sendJson(response, 404, { error: "not found" });
}

function listen(listener, host) {
	const server = createServer((request, response) => {
		handle(listener, request, response).catch((error) => {
			log("error", "request failed", { listener, error: error.message });
			sendJson(response, 500, { ok: false, error: error.message });
		});
	});
	server.on("error", (error) => {
		// Most likely the private address is not up yet. Exit and let the supervisor retry, rather than
		// run half a broker that the other hosts cannot reach.
		log("error", "listen failed", { listener, host, port: PORT, error: error.message });
		process.exit(1);
	});
	server.listen(PORT, host, () => log("info", "listening", { listener, host, port: PORT, version: readVersion() }));
	return server;
}

const servers = [listen("ext", EXT_HOST)];
if (PEER_HOST && PEER_HOST !== EXT_HOST) servers.push(listen("peer", PEER_HOST));

function shutdown(signal) {
	log("info", "shutting down", { signal, pending: pending.size });
	for (const entry of pending.values()) {
		clearTimeout(entry.timer);
		entry.reject(new Error("pi-chrome broker stopped"));
	}
	pending.clear();
	for (const waiter of [...waiters]) waiter(undefined);
	for (const server of servers) server.close();
	setTimeout(() => process.exit(0), 1_000).unref();
}
process.on("SIGTERM", () => shutdown("SIGTERM"));
process.on("SIGINT", () => shutdown("SIGINT"));