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

/**
 * SSRF guard for user-supplied endpoint URLs. The PolyLLM server fetches these URLs with the user's
 * headers, so without a guard a user could probe the server's own network (cloud metadata, Redis,
 * the database…). Private, loopback, link-local and special-purpose ranges are rejected unless the
 * operator sets `ALLOW_PRIVATE_ENDPOINTS=1` (self-hosted PolyLLM next to Ollama / LM Studio).
 *
 * Pure helpers (`isPrivateIp`, `isBlockedHostname`, `checkEndpointUrlSync`) are unit-tested;
 * `assertEndpointUrlAllowed` additionally resolves the hostname so `evil.example → 10.0.0.1` is caught.
 */

export const ALLOWED_PROTOCOLS = new Set(["http:", "https:"]);

export function privateEndpointsAllowed(env: NodeJS.ProcessEnv = process.env): boolean {
  return env.ALLOW_PRIVATE_ENDPOINTS === "1" || env.ALLOW_PRIVATE_ENDPOINTS === "true";
}

function ipv4ToInt(ip: string): number | null {
  const parts = ip.split(".");
  if (parts.length !== 4) return null;
  let n = 0;
  for (const p of parts) {
    if (!/^\d{1,3}$/.test(p)) return null;
    const v = Number(p);
    if (v > 255) return null;
    n = n * 256 + v;
  }
  return n;
}

const V4_BLOCKS: [string, number][] = [
  ["0.0.0.0", 8], // "this" network
  ["10.0.0.0", 8], // private
  ["100.64.0.0", 10], // carrier-grade NAT
  ["127.0.0.0", 8], // loopback
  ["169.254.0.0", 16], // link-local (cloud metadata lives here)
  ["172.16.0.0", 12], // private
  ["192.0.0.0", 24], // IETF protocol assignments
  ["192.0.2.0", 24], // TEST-NET-1
  ["192.168.0.0", 16], // private
  ["198.18.0.0", 15], // benchmarking
  ["198.51.100.0", 24], // TEST-NET-2
  ["203.0.113.0", 24], // TEST-NET-3
  ["224.0.0.0", 4], // multicast
  ["240.0.0.0", 4], // reserved + broadcast
];

function inV4Block(ip: number, base: string, bits: number): boolean {
  const b = ipv4ToInt(base)!;
  const mask = bits === 0 ? 0 : (~0 << (32 - bits)) >>> 0;
  return ((ip & mask) >>> 0) === ((b & mask) >>> 0);
}

export function isPrivateIpv4(ip: string): boolean {
  const n = ipv4ToInt(ip);
  if (n === null) return false;
  return V4_BLOCKS.some(([base, bits]) => inV4Block(n, base, bits));
}

/** Expand an IPv6 literal into 8 hextets (handles `::` and embedded IPv4). Returns null when malformed. */
function expandIpv6(ip: string): number[] | null {
  let s = ip.toLowerCase();
  const zone = s.indexOf("%");
  if (zone >= 0) s = s.slice(0, zone);
  // embedded IPv4 tail → two hextets
  const lastColon = s.lastIndexOf(":");
  if (s.includes(".") && lastColon >= 0) {
    const v4 = ipv4ToInt(s.slice(lastColon + 1));
    if (v4 === null) return null;
    s = `${s.slice(0, lastColon)}:${((v4 >>> 16) & 0xffff).toString(16)}:${(v4 & 0xffff).toString(16)}`;
  }
  const halves = s.split("::");
  if (halves.length > 2) return null;
  const head = halves[0] ? halves[0].split(":") : [];
  const tail = halves.length === 2 && halves[1] ? halves[1].split(":") : [];
  const fill = halves.length === 2 ? 8 - head.length - tail.length : 0;
  if (fill < 0 || (halves.length === 1 && head.length !== 8)) return null;
  const parts = [...head, ...Array<string>(fill).fill("0"), ...tail];
  if (parts.length !== 8) return null;
  const out: number[] = [];
  for (const p of parts) {
    if (!/^[0-9a-f]{1,4}$/.test(p)) return null;
    out.push(parseInt(p, 16));
  }
  return out;
}

export function isPrivateIpv6(ip: string): boolean {
  const h = expandIpv6(ip);
  if (!h) return false;
  const allZero = h.every((x) => x === 0);
  if (allZero) return true; // ::
  if (h.slice(0, 7).every((x) => x === 0) && h[7] === 1) return true; // ::1
  // IPv4-mapped (::ffff:a.b.c.d) and IPv4-compatible → check the embedded v4
  if (h.slice(0, 5).every((x) => x === 0) && (h[5] === 0xffff || h[5] === 0)) {
    const v4 = `${h[6] >> 8}.${h[6] & 0xff}.${h[7] >> 8}.${h[7] & 0xff}`;
    if (h[5] === 0xffff || h[6] !== 0 || h[7] > 1) return isPrivateIpv4(v4);
  }
  // 64:ff9b::/96 (NAT64) → embedded v4
  if (h[0] === 0x64 && h[1] === 0xff9b && h.slice(2, 6).every((x) => x === 0)) return isPrivateIpv4(`${h[6] >> 8}.${h[6] & 0xff}.${h[7] >> 8}.${h[7] & 0xff}`);
  if ((h[0] & 0xfe00) === 0xfc00) return true; // fc00::/7 unique local
  if ((h[0] & 0xffc0) === 0xfe80) return true; // fe80::/10 link-local
  if ((h[0] & 0xff00) === 0xff00) return true; // multicast
  if (h[0] === 0x2001 && h[1] === 0x0db8) return true; // documentation
  return false;
}

export function isPrivateIp(ip: string): boolean {
  const v = isIP(ip);
  if (v === 4) return isPrivateIpv4(ip);
  if (v === 6) return isPrivateIpv6(ip);
  return false;
}

/** Hostnames that always mean "this machine / this network" without needing DNS. */
export function isBlockedHostname(hostname: string): boolean {
  const h = hostname.toLowerCase().replace(/\.$/, "");
  if (h === "localhost" || h.endsWith(".localhost")) return true;
  if (h.endsWith(".local") || h.endsWith(".internal") || h.endsWith(".lan") || h.endsWith(".home") || h.endsWith(".home.arpa") || h.endsWith(".corp") || h.endsWith(".intranet")) return true;
  if (h === "metadata.google.internal" || h === "metadata") return true;
  if (!h.includes(".") && isIP(h) === 0) return true; // bare single-label names resolve inside the server's search domain
  return false;
}

export interface EndpointUrlCheck {
  ok: boolean;
  /** Stable machine-readable reason. */
  reason?: "INVALID_URL" | "BAD_PROTOCOL" | "CREDENTIALS_IN_URL" | "PRIVATE_HOST" | "PRIVATE_IP" | "PRIVATE_DNS";
  message?: string;
  /** True when the URL points at a private/loopback host (allowed or not). */
  isPrivate: boolean;
  url?: URL;
}

/** Synchronous part of the check (no DNS). */
export function checkEndpointUrlSync(raw: string, opts: { allowPrivate?: boolean } = {}): EndpointUrlCheck {
  let url: URL;
  try {
    url = new URL(raw.trim());
  } catch {
    return { ok: false, reason: "INVALID_URL", message: "Enter a full URL such as http://localhost:11434/v1.", isPrivate: false };
  }
  if (!ALLOWED_PROTOCOLS.has(url.protocol)) return { ok: false, reason: "BAD_PROTOCOL", message: "Only http:// and https:// endpoints are supported.", isPrivate: false, url };
  if (url.username || url.password) return { ok: false, reason: "CREDENTIALS_IN_URL", message: "Put credentials in the API key or headers, not in the URL.", isPrivate: false, url };
  const host = url.hostname.replace(/^\[|\]$/g, "");
  const allowPrivate = opts.allowPrivate ?? privateEndpointsAllowed();
  const literalIp = isIP(host) !== 0;
  const isPrivate = literalIp ? isPrivateIp(host) : isBlockedHostname(host);
  if (isPrivate && !allowPrivate) {
    return {
      ok: false,
      reason: literalIp ? "PRIVATE_IP" : "PRIVATE_HOST",
      message: "This address is on a private or local network the PolyLLM server cannot reach. Expose it through a tunnel (Cloudflare, ngrok, Tailscale Funnel…) or run PolyLLM locally with ALLOW_PRIVATE_ENDPOINTS=1.",
      isPrivate,
      url,
    };
  }
  return { ok: true, isPrivate, url };
}

/** Full check: syntax + DNS resolution of the hostname against private ranges (DNS rebinding is out of scope). */
export async function assertEndpointUrlAllowed(raw: string, opts: { allowPrivate?: boolean } = {}): Promise<EndpointUrlCheck> {
  const sync = checkEndpointUrlSync(raw, opts);
  if (!sync.ok || !sync.url) return sync;
  const allowPrivate = opts.allowPrivate ?? privateEndpointsAllowed();
  if (allowPrivate || sync.isPrivate) return sync;
  const host = sync.url.hostname.replace(/^\[|\]$/g, "");
  if (isIP(host)) return sync;
  try {
    const addrs = await lookup(host, { all: true, verbatim: true });
    if (addrs.some((a) => isPrivateIp(a.address))) {
      return { ...sync, ok: false, reason: "PRIVATE_DNS", isPrivate: true, message: `${host} resolves to a private address the PolyLLM server will not call.` };
    }
  } catch {
    // Unresolvable hosts are reported by the connection test itself, not here.
  }
  return sync;
}