import { lookup as dnsLookup } from "node:dns/promises";
import net from "node:net";

export type LookupAddress = { address: string; family: number };
export type DnsLookup = (hostname: string) => Promise<LookupAddress[]>;

async function defaultLookup(hostname: string): Promise<LookupAddress[]> {
  return dnsLookup(hostname, { all: true, verbatim: true });
}

function normalizeHostname(hostname: string): string {
  return hostname.toLowerCase().replace(/^\[|\]$/g, "").replace(/\.$/, "");
}

function blockedIPv4(address: string): boolean {
  const octets = address.split(".").map(Number);
  if (octets.length !== 4 || octets.some((part) => !Number.isInteger(part) || part < 0 || part > 255)) return true;
  const [a = 0, b = 0, c = 0] = octets;
  return (
    a === 0 ||
    a === 10 ||
    a === 127 ||
    (a === 100 && b >= 64 && b <= 127) ||
    (a === 169 && b === 254) ||
    (a === 172 && b >= 16 && b <= 31) ||
    (a === 192 && b === 0 && c === 0) ||
    (a === 192 && b === 0 && c === 2) ||
    (a === 192 && b === 168) ||
    (a === 198 && (b === 18 || b === 19)) ||
    (a === 198 && b === 51 && c === 100) ||
    (a === 203 && b === 0 && c === 113) ||
    a >= 224
  );
}

function expandIPv6(address: string): number[] | undefined {
  let normalized = address.toLowerCase().split("%")[0] ?? "";
  if (normalized.includes(".")) {
    const lastColon = normalized.lastIndexOf(":");
    const ipv4 = normalized.slice(lastColon + 1);
    if (net.isIP(ipv4) !== 4) return undefined;
    const octets = ipv4.split(".").map(Number);
    normalized = `${normalized.slice(0, lastColon)}:${(((octets[0] ?? 0) << 8) | (octets[1] ?? 0)).toString(16)}:${(((octets[2] ?? 0) << 8) | (octets[3] ?? 0)).toString(16)}`;
  }
  const halves = normalized.split("::");
  if (halves.length > 2) return undefined;
  const left = halves[0] ? halves[0].split(":") : [];
  const right = halves.length === 2 && halves[1] ? halves[1].split(":") : [];
  const missing = 8 - left.length - right.length;
  if ((halves.length === 1 && missing !== 0) || missing < 0) return undefined;
  const groups = [...left, ...Array(missing).fill("0"), ...right].map((part) => {
    if (!/^[0-9a-f]{1,4}$/.test(part)) return -1;
    return Number.parseInt(part, 16);
  });
  return groups.length === 8 && groups.every((part) => part >= 0 && part <= 0xffff) ? groups : undefined;
}

function blockedIPv6(address: string): boolean {
  const groups = expandIPv6(address);
  if (!groups) return true;
  const first = groups[0] ?? 0;
  if (groups.every((group) => group === 0)) return true;
  if (groups.slice(0, 7).every((group) => group === 0) && groups[7] === 1) return true;
  if ((first & 0xfe00) === 0xfc00) return true;
  if ((first & 0xffc0) === 0xfe80) return true;
  if ((first & 0xff00) === 0xff00) return true;
  if (first === 0x2001 && groups[1] === 0x0db8) return true;
  const mapped = groups.slice(0, 5).every((group) => group === 0) && groups[5] === 0xffff;
  if (mapped) {
    const sixth = groups[6] ?? 0;
    const seventh = groups[7] ?? 0;
    return blockedIPv4(`${sixth >> 8}.${sixth & 0xff}.${seventh >> 8}.${seventh & 0xff}`);
  }
  return false;
}

export function assertPublicAddress(address: string, hostname: string): void {
  const normalized = normalizeHostname(address);
  const family = net.isIP(normalized);
  if (family === 0) throw new Error(`Resolved non-IP address for ${hostname}: ${address}`);
  if ((family === 4 && blockedIPv4(normalized)) || (family === 6 && blockedIPv6(normalized))) {
    throw new Error(`Blocked non-public address for ${hostname}: ${normalized}`);
  }
}

export async function validatePublicHttpUrl(rawUrl: string | URL, lookup: DnsLookup = defaultLookup): Promise<URL> {
  const url = rawUrl instanceof URL ? new URL(rawUrl) : new URL(rawUrl);
  if (url.protocol !== "http:" && url.protocol !== "https:") throw new Error("web_fetch only supports HTTP and HTTPS URLs");
  if (url.username || url.password) throw new Error("web_fetch does not allow credentials in URLs");
  const hostname = normalizeHostname(url.hostname);
  if (!hostname) throw new Error("URL must include a hostname");
  if (hostname === "localhost" || hostname.endsWith(".localhost") || hostname.endsWith(".local")) {
    throw new Error(`Blocked local hostname: ${hostname}`);
  }
  if (net.isIP(hostname)) {
    assertPublicAddress(hostname, hostname);
    return url;
  }
  let addresses: LookupAddress[];
  try {
    addresses = await lookup(hostname);
  } catch (error) {
    throw new Error(`Failed to resolve ${hostname}: ${error instanceof Error ? error.message : String(error)}`);
  }
  if (addresses.length === 0) throw new Error(`Failed to resolve ${hostname}: no addresses returned`);
  for (const entry of addresses) assertPublicAddress(entry.address, hostname);
  return url;
}