import { lookup } from "node:dns/promises";
import { isIP } from "node:net";
import * as CID from "@atcute/cid";
import { isAuthEnabled } from "../atproto/env.ts";
import { AppError } from "./errors.ts";
import { resolvePdsEndpoint } from "./identity.ts";
import { LIMITS } from "./limits.ts";

type BlobFetchReason =
	| "invalid-cid"
	| "pds-resolve-failed"
	| "pds-not-found"
	| "non-https-pds"
	| "blocked-host"
	| "fetch-failed"
	| "timeout"
	| "upstream-error"
	| "too-large"
	| "cid-mismatch";

export class BlobFetchError extends AppError {
	readonly reason: BlobFetchReason;
	constructor(reason: BlobFetchReason, status: number, message: string) {
		super(message, status);
		this.reason = reason;
	}
}

interface VerifiedBlob {
	data: Uint8Array;
	mimeType: string;
}

/**
 * Private, loopback, link-local, ULA, and CGNAT ranges — the addresses an SSRF
 * payload would aim at to reach internal services or the cloud metadata endpoint
 * (169.254.169.254).
 */
export function isPrivateIp(ip: string): boolean {
	if (isIP(ip) === 4) {
		const parts = ip.split(".");
		const a = Number(parts[0]);
		const b = Number(parts[1]);
		return (
			a === 0 ||
			a === 10 ||
			a === 127 ||
			(a === 169 && b === 254) || // link-local incl. cloud metadata
			(a === 172 && b >= 16 && b <= 31) ||
			(a === 192 && b === 168) ||
			(a === 100 && b >= 64 && b <= 127) // CGNAT
		);
	}
	const v6 = ip.toLowerCase();
	if (v6 === "::1" || v6 === "::") return true;
	if (v6.startsWith("fe80") || v6.startsWith("fc") || v6.startsWith("fd")) {
		return true; // link-local + unique-local
	}
	if (v6.startsWith("::ffff:")) return isPrivateIp(v6.slice(7)); // mapped IPv4
	return false;
}

/**
 * SSRF guard for the resolved PDS host. A DID document can point its
 * serviceEndpoint at any host (e.g. via did:web), so before connecting we reject
 * hosts that resolve to private/reserved addresses. `allowLocal` (dev only)
 * turns the guard off so a local PDS works. Throws a generic BlobFetchError so
 * the caller can't be used as a probe of internal network reachability.
 */
export async function assertSafePdsHost(
	hostname: string,
	allowLocal: boolean,
): Promise<void> {
	if (allowLocal) return;
	if (isIP(hostname)) {
		if (isPrivateIp(hostname)) {
			throw new BlobFetchError("blocked-host", 502, "Blob not found");
		}
		return;
	}
	const addrs = await lookup(hostname, { all: true });
	for (const { address } of addrs) {
		if (isPrivateIp(address)) {
			throw new BlobFetchError("blocked-host", 502, "Blob not found");
		}
	}
}

/**
 * Fetch a blob from its owning PDS and verify it.
 *
 * The PDS is resolved from the DID, then `com.atproto.sync.getBlob` is called
 * with these guarantees:
 *
 *   - SSRF guard: only HTTPS endpoints, except localhost for dev
 *   - Timeout: aborts after `LIMITS.blobProxy.timeoutMs`
 *   - Size cap: rejects on Content-Length, then enforces during streaming
 *     (the streaming check is the actual security boundary — a malicious PDS
 *     can lie about Content-Length)
 *   - CID verification: hashes the bytes and rejects if they don't match the
 *     requested CID (a malicious PDS could otherwise serve arbitrary content
 *     for any CID)
 */
export async function fetchVerifiedBlob(
	did: string,
	cid: string,
	fetchFn: typeof fetch = fetch,
	resolveEndpoint: (did: string) => Promise<string | null> = resolvePdsEndpoint,
): Promise<VerifiedBlob> {
	let expected: CID.Cid;
	try {
		expected = CID.fromString(cid);
	} catch {
		throw new BlobFetchError("invalid-cid", 400, "Invalid CID");
	}

	let pdsEndpoint: string;
	try {
		const resolved = await resolveEndpoint(did);
		if (!resolved) {
			throw new BlobFetchError("pds-not-found", 404, "PDS not found for DID");
		}
		pdsEndpoint = resolved;
	} catch (err) {
		if (err instanceof BlobFetchError) throw err;
		throw new BlobFetchError(
			"pds-resolve-failed",
			502,
			"Failed to resolve DID",
		);
	}

	// localhost over http is a dev-only convenience; in prod every PDS is HTTPS.
	const allowLocal = !isAuthEnabled();
	const pdsUrl = new URL(pdsEndpoint);
	const isLocalhost =
		pdsUrl.hostname === "localhost" || pdsUrl.hostname === "127.0.0.1";
	if (pdsUrl.protocol !== "https:" && !(allowLocal && isLocalhost)) {
		throw new BlobFetchError(
			"non-https-pds",
			502,
			"PDS endpoint must be HTTPS",
		);
	}
	await assertSafePdsHost(pdsUrl.hostname, allowLocal);

	const blobUrl = `${pdsEndpoint}/xrpc/com.atproto.sync.getBlob?did=${encodeURIComponent(did)}&cid=${encodeURIComponent(cid)}`;

	let upstream: Response;
	try {
		upstream = await fetchFn(blobUrl, {
			signal: AbortSignal.timeout(LIMITS.blobProxy.timeoutMs),
		});
	} catch (err) {
		// Normalize timeout vs connection-refused vs other to one generic 502 so
		// the proxy can't be probed for internal port/host reachability. (Timing
		// still leaks somewhat; the per-IP rate limit caps that.)
		if (err instanceof Error && err.name === "TimeoutError") {
			throw new BlobFetchError("timeout", 502, "Blob not found");
		}
		throw new BlobFetchError("fetch-failed", 502, "Blob not found");
	}

	if (!upstream.ok) {
		// Was `upstream.status`: passing the upstream status through leaked it as
		// an oracle. Collapse every non-2xx to the same generic 502.
		await upstream.body?.cancel().catch(() => {});
		throw new BlobFetchError("upstream-error", 502, "Blob not found");
	}

	const announced = upstream.headers.get("content-length");
	if (announced !== null && Number(announced) > LIMITS.blobProxy.maxBytes) {
		await upstream.body?.cancel().catch(() => {});
		throw new BlobFetchError("too-large", 413, "Blob too large");
	}

	const mimeType =
		upstream.headers.get("content-type") ?? "application/octet-stream";

	if (upstream.body === null) {
		throw new BlobFetchError("upstream-error", 502, "Blob not found");
	}

	const reader = upstream.body.getReader();
	const chunks: Uint8Array[] = [];
	let total = 0;
	try {
		while (true) {
			const { done, value } = await reader.read();
			if (done) break;
			total += value.byteLength;
			// Real size boundary: a malicious PDS can lie about Content-Length.
			if (total > LIMITS.blobProxy.maxBytes) {
				await reader.cancel().catch(() => {});
				throw new BlobFetchError("too-large", 413, "Blob too large");
			}
			chunks.push(value);
		}
	} catch (err) {
		if (err instanceof BlobFetchError) throw err;
		throw new BlobFetchError("fetch-failed", 502, "Blob not found");
	}

	const data = new Uint8Array(total);
	let offset = 0;
	for (const chunk of chunks) {
		data.set(chunk, offset);
		offset += chunk.byteLength;
	}

	// Verify bytes hash to the requested CID, else a PDS could serve anything for any CID.
	const computed = await CID.create(0x55, data);
	if (!CID.equals(computed, expected)) {
		throw new BlobFetchError(
			"cid-mismatch",
			502,
			"Blob content does not match CID",
		);
	}

	return { data, mimeType };
}