import { HttpError, badRequest } from "../errors";

const DOH_ENDPOINT = "https://cloudflare-dns.com/dns-query";
const DOH_TIMEOUT_MS = 3_000;
const HTTP_UNAVAILABLE = 503;
const DNS_NOERROR = 0;
const DNS_NXDOMAIN = 3;
const CACHE_TTL_MS = 60_000;
const MAX_CACHE_ENTRIES = 500;
const MAX_CNAME_DEPTH = 5;
const MAX_DOH_BYTES = 64 * 1024;
const TYPE_A = 1;
const TYPE_CNAME = 5;
const TYPE_AAAA = 28;
const IPV4_OCTETS = 4;
const IPV6_GROUPS = 8;
const BYTE = 256;
const GROUP = 0x10000;

type Answer = { type: number; data: string };
type DohResponse = { Status: number; Answer?: Answer[] };
export type Resolver = (name: string, type: "A" | "AAAA") => Promise<DohResponse>;

/** Non-public IPv4 space as CIDRs (RFC 6890 special-purpose registry plus multicast and reserved). */
const PRIVATE_V4_CIDRS: Array<[string, number]> = [
  ["0.0.0.0", 8],
  ["10.0.0.0", 8],
  ["100.64.0.0", 10],
  ["127.0.0.0", 8],
  ["169.254.0.0", 16],
  ["172.16.0.0", 12],
  ["192.0.0.0", 24],
  ["192.0.2.0", 24],
  ["192.88.99.0", 24],
  ["192.168.0.0", 16],
  ["198.18.0.0", 15],
  ["198.51.100.0", 24],
  ["203.0.113.0", 24],
  ["224.0.0.0", 4],
  ["240.0.0.0", 4],
];
const IPV4_BITS = 32;

function ipv4ToInt(octets: number[]): number {
  return octets.reduce((value, octet) => value * BYTE + octet, 0);
}

const PRIVATE_V4_RANGES = PRIVATE_V4_CIDRS.map(([base, bits]) => {
  const size = 2 ** (IPV4_BITS - bits);
  const start = ipv4ToInt(base.split(".").map(Number));
  return [start, start + size - 1] as const;
});

function parseIpv4(text: string): number[] | null {
  const parts = text.split(".");
  if (parts.length !== IPV4_OCTETS) return null;
  const octets = parts.map((part) => (/^\d{1,3}$/.test(part) ? Number(part) : Number.NaN));
  return octets.every((octet) => octet >= 0 && octet < BYTE) ? octets : null;
}

function isPrivateIpv4(octets: number[]): boolean {
  const value = ipv4ToInt(octets);
  return PRIVATE_V4_RANGES.some(([start, end]) => value >= start && value <= end);
}

/** Expands an IPv6 literal (with optional embedded IPv4 tail) into 8 16-bit groups. */
function parseIpv6(text: string): number[] | null {
  let value = text
    .toLowerCase()
    .replace(/^\[|\]$/g, "")
    .split("%")[0];
  const tail = /(\d{1,3}(?:\.\d{1,3}){3})$/.exec(value);
  if (tail) {
    const octets = parseIpv4(tail[1]);
    if (!octets) return null;
    value = `${value.slice(0, -tail[1].length)}${((octets[0] << 8) | octets[1]).toString(16)}:${((octets[2] << 8) | octets[3]).toString(16)}`;
  }
  const halves = value.split("::");
  if (halves.length > 2) return null;
  const head = halves[0] ? halves[0].split(":") : [];
  const rest = halves.length === 2 && halves[1] ? halves[1].split(":") : [];
  const missing = IPV6_GROUPS - head.length - rest.length;
  if (halves.length === 1 ? missing !== 0 : missing < 1) return null;
  const groups = [...head, ...Array<string>(halves.length === 2 ? missing : 0).fill("0"), ...rest].map((group) =>
    /^[0-9a-f]{1,4}$/.test(group) ? Number.parseInt(group, 16) : Number.NaN,
  );
  return groups.every((group) => group >= 0 && group < GROUP) ? groups : null;
}

function embeddedIpv4(groups: number[]): number[] {
  return [...groupToOctets(groups[6]), ...groupToOctets(groups[7])];
}

/** [mask, value] on the first 16-bit group: unique-local fc00::/7, link-local fe80::/10, multicast ff00::/8. */
const PRIVATE_V6_PREFIXES: Array<[number, number]> = [
  [0xfe00, 0xfc00],
  [0xffc0, 0xfe80],
  [0xffc0, 0xfec0],
  [0xff00, 0xff00],
];
/** Only 2000::/3 is global unicast; 2001::/23 is IETF protocol space (Teredo, benchmarking, ORCHID…). */
const GLOBAL_UNICAST_MASK = 0xe000;
const GLOBAL_UNICAST = 0x2000;
const IETF_V6 = [0x2001, 0x0200];
const V4_MAPPED = 0xffff;
const NAT64 = [0x64, 0xff9b];
const SIX_TO_FOUR = 0x2002;
const DOCUMENTATION = [0x2001, 0x0db8];
const ZERO_PREFIX_GROUPS = 5;

function groupToOctets(group: number): number[] {
  return [group >> 8, group & 0xff];
}

function isPrivateIpv6(groups: number[]): boolean {
  const zeroPrefix = groups.slice(0, ZERO_PREFIX_GROUPS).every((group) => group === 0);
  const [, , , , , mapped, high, low] = groups;
  if (zeroPrefix && mapped === 0 && high === 0 && low <= 1) return true;
  if (zeroPrefix && (mapped === V4_MAPPED || mapped === 0)) return isPrivateIpv4(embeddedIpv4(groups));
  if (groups.slice(0, 4).every((group) => group === 0) && groups[4] === V4_MAPPED && mapped === 0) {
    return isPrivateIpv4(embeddedIpv4(groups));
  }
  if (groups[0] === NAT64[0] && groups[1] === NAT64[1]) {
    const wellKnown = groups.slice(2, 6).every((group) => group === 0);
    return !wellKnown || isPrivateIpv4(embeddedIpv4(groups));
  }
  if ((groups[0] & GLOBAL_UNICAST_MASK) !== GLOBAL_UNICAST) return true;
  if (groups[0] === IETF_V6[0] && groups[1] < IETF_V6[1]) return true;
  if (groups[0] === SIX_TO_FOUR) return isPrivateIpv4([...groupToOctets(groups[1]), ...groupToOctets(groups[2])]);
  if (groups[0] === DOCUMENTATION[0] && groups[1] === DOCUMENTATION[1]) return true;
  return PRIVATE_V6_PREFIXES.some(([mask, value]) => (groups[0] & mask) === value);
}

/** True for loopback, private, link-local, CGNAT, metadata, multicast and reserved addresses (v4 and v6). */
export function isNonPublicAddress(address: string): boolean {
  const v4 = parseIpv4(address);
  if (v4) return isPrivateIpv4(v4);
  const v6 = parseIpv6(address);
  return v6 ? isPrivateIpv6(v6) : true;
}

/**
 * For a URL `hostname` (as produced by `new URL()`, so IPv4 shorthand such as `127.1` or `0x7f.1` is already
 * canonical and IPv6 is bracketed): `true`/`false` when it is an IP literal, `null` when it is a DNS name.
 * Callers that only have a hostname filter use this instead of hand-rolled prefix regexes (CI-A9).
 */
export function isNonPublicIpLiteral(hostname: string): boolean | null {
  const host = hostname.toLowerCase().replace(/\.$/, "");
  const v4 = parseIpv4(host);
  if (v4) return isPrivateIpv4(v4);
  if (host.startsWith("[") || host.includes(":")) return isNonPublicAddress(host);
  // Numeric forms that are not canonical dotted-decimal (`127.1`, `0x7f.1`, `0177.0.0.1`, `2130706433`) are IPv4
  // shorthands to libc resolvers; refuse them instead of treating them as DNS names.
  if (/^(0x[0-9a-f]+|\d+)(\.(0x[0-9a-f]+|\d+))*$/.test(host)) return true;
  return null;
}

/** DoH outages (timeouts, 5xx, SERVFAIL) are retryable and distinct from a name that does not exist. */
export function unverifiableAddress(): HttpError {
  return new HttpError(HTTP_UNAVAILABLE, "Could not verify the source address; try again");
}

const dohResolver: Resolver = async (name, type) => {
  const url = `${DOH_ENDPOINT}?name=${encodeURIComponent(name)}&type=${type}`;
  let parsed: DohResponse;
  try {
    const response = await fetch(url, {
      headers: { accept: "application/dns-json" },
      redirect: "manual",
      signal: AbortSignal.timeout(DOH_TIMEOUT_MS),
    });
    if (!response.ok || Number(response.headers.get("content-length") ?? "0") > MAX_DOH_BYTES) {
      await response.body?.cancel();
      throw unverifiableAddress();
    }
    const text = await response.text();
    if (text.length > MAX_DOH_BYTES) throw unverifiableAddress();
    parsed = JSON.parse(text) as DohResponse;
  } catch {
    throw unverifiableAddress();
  }
  if (parsed.Status !== DNS_NOERROR && parsed.Status !== DNS_NXDOMAIN) throw unverifiableAddress();
  return parsed;
};

/** Least-recently-used map: a flood of unique hosts evicts the oldest entries instead of clearing hot ones (GIT-P10). */
export class LruCache<V> {
  private readonly entries = new Map<string, V>();

  constructor(private readonly capacity: number) {}

  get(key: string): V | undefined {
    const value = this.entries.get(key);
    if (value === undefined) return undefined;
    this.entries.delete(key);
    this.entries.set(key, value);
    return value;
  }

  set(key: string, value: V): void {
    this.entries.delete(key);
    this.entries.set(key, value);
    while (this.entries.size > this.capacity) {
      const oldest = this.entries.keys().next();
      if (oldest.done) break;
      this.entries.delete(oldest.value);
    }
  }

  get size(): number {
    return this.entries.size;
  }
}

const cache = new LruCache<{ addresses: string[]; expires: number }>(MAX_CACHE_ENTRIES);

async function lookup(name: string, resolver: Resolver, depth = 0): Promise<string[]> {
  if (depth > MAX_CNAME_DEPTH) throw badRequest(`Too many CNAMEs for ${name}`);
  const results = await Promise.all([resolver(name, "A"), resolver(name, "AAAA")]);
  const addresses: string[] = [];
  const aliases = new Set<string>();
  for (const result of results) {
    for (const answer of result.Answer ?? []) {
      if (answer.type === TYPE_A || answer.type === TYPE_AAAA) addresses.push(answer.data);
      else if (answer.type === TYPE_CNAME) aliases.add(answer.data.replace(/\.$/, "").toLowerCase());
    }
  }
  if (addresses.length === 0) {
    for (const alias of aliases) addresses.push(...(await lookup(alias, resolver, depth + 1)));
  }
  return addresses;
}

/**
 * Resolves the URL's host over DNS-over-HTTPS (A and AAAA, following CNAMEs) and refuses it when any answer
 * is non-public, or when there is none. DoH outages throw a retryable 503 so callers fail closed but clearly.
 *
 * Answers (including "does not resolve") are cached per host for 60s per isolate.
 *
 * TOCTOU: Workers' fetch() cannot pin a connection to the address checked here (no resolve override or
 * connect-to-IP with SNI), so a DNS-rebinding window remains between this lookup and Cloudflare's own
 * resolution inside fetch(). Mitigations: answers are cached for 60s per isolate so repeated checks agree,
 * callers never follow redirects without re-checking each hop, and Workers egress cannot reach RFC 1918,
 * loopback or cloud-metadata addresses on Cloudflare's network in the first place. This residual window is an
 * accepted risk (GIT-P9).
 */
export async function assertResolvesPublic(raw: string | URL, resolver: Resolver = dohResolver): Promise<void> {
  const host = (typeof raw === "string" ? new URL(raw) : raw).hostname.toLowerCase().replace(/\.$/, "");
  if (parseIpv4(host) || host.startsWith("[")) {
    if (isNonPublicAddress(host)) throw badRequest("The URL points at a non-public address");
    return;
  }
  const now = Date.now();
  let addresses = resolver === dohResolver ? cache.get(host) : undefined;
  if (!addresses || addresses.expires < now) {
    const resolved = await lookup(host, resolver);
    addresses = { addresses: resolved, expires: now + CACHE_TTL_MS };
    if (resolver === dohResolver) cache.set(host, addresses);
  }
  if (addresses.addresses.length === 0) throw badRequest(`${host} does not resolve`);
  if (addresses.addresses.some(isNonPublicAddress)) throw badRequest(`${host} resolves to a non-public address`);
}