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

const PRIVATE_V4_CIDRS = [
  ["127.0.0.0", 8],
  ["10.0.0.0", 8],
  ["172.16.0.0", 12],
  ["192.168.0.0", 16],
  ["169.254.0.0", 16],
  ["0.0.0.0", 8],
  ["100.64.0.0", 10],
] as const;

function ipv4ToInt(ip: string): number {
  const parts = ip.split(".").map((n) => Number(n));
  if (parts.length !== 4 || parts.some((n) => Number.isNaN(n) || n < 0 || n > 255)) return -1;
  return (((parts[0]! << 24) | (parts[1]! << 16) | (parts[2]! << 8) | parts[3]!) >>> 0);
}

function isPrivateV4(ip: string): boolean {
  const target = ipv4ToInt(ip);
  if (target < 0) return true;
  for (const [base, bits] of PRIVATE_V4_CIDRS) {
    const baseInt = ipv4ToInt(base);
    const mask = bits === 0 ? 0 : (~0 << (32 - bits)) >>> 0;
    if ((target & mask) === (baseInt & mask)) return true;
  }
  return false;
}

function isPrivateV6(ip: string): boolean {
  const lower = ip.toLowerCase();
  if (lower === "::1" || lower === "::") return true;
  if (lower.startsWith("fc") || lower.startsWith("fd")) return true; // fc00::/7
  if (lower.startsWith("fe80")) return true; // link-local
  if (lower.startsWith("::ffff:")) {
    // IPv4-mapped
    const v4 = lower.slice("::ffff:".length);
    if (net.isIPv4(v4)) return isPrivateV4(v4);
  }
  return false;
}

/**
 * Resolve a hostname and return true if any A/AAAA record is a private/loopback address.
 * Used to prevent SSRF when an LLM crafts a URL pointing to internal services.
 */
export async function isPrivateHost(hostname: string): Promise<boolean> {
  const host = hostname.toLowerCase();
  if (host === "localhost" || host.endsWith(".localhost")) return true;
  if (net.isIPv4(host)) return isPrivateV4(host);
  if (net.isIPv6(host)) return isPrivateV6(host);
  try {
    const addrs = await dns.lookup(host, { all: true });
    for (const a of addrs) {
      if (a.family === 4 && isPrivateV4(a.address)) return true;
      if (a.family === 6 && isPrivateV6(a.address)) return true;
    }
  } catch {
    // Unresolvable: treat as suspicious
    return true;
  }
  return false;
}