/**
 * SSRF-guarded fetch + bounded body reads.
 *
 * Page fetching is steerable by tool arguments and (transitively) by injected
 * instructions in fetched pages, so every outbound request is gated:
 *   - scheme allowlist (http/https only),
 *   - DNS-resolve the hostname and reject loopback/private/link-local/ULA/etc.
 *     addresses (blocks cloud-metadata 169.254.169.254, localhost, RFC-1918…),
 *   - manual redirect handling that re-validates EVERY hop (a public URL can
 *     302 to an internal one), capped to bound the chain,
 *   - streamed body reads with a hard byte cap so a hostile/huge response can't
 *     exhaust memory.
 *
 * `allowPrivateNetwork` (global config only) lifts the address check for users
 * who intentionally fetch localhost/internal docs.
 *
 * The IP classifier (`isBlockedAddress`) is pure and exported so it can be
 * unit-tested without touching the network.
 *
 * Caveat: resolve-then-connect is best-effort against DNS rebinding (the kernel
 * re-resolves on connect). It blocks the overwhelming majority of SSRF attempts
 * but is not a substitute for an egress firewall in hostile environments.
 */

import dns from "node:dns/promises";
import net from "node:net";

// ---------------------------------------------------------------------------
// IP classification (pure)
// ---------------------------------------------------------------------------

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

function isBlockedIPv4(ip: string): boolean {
  const n = ipv4ToInt(ip);
  if (n === undefined) return true; // unparseable -> treat as unsafe
  const inRange = (base: string, bits: number) => {
    const b = ipv4ToInt(base)!;
    const mask = bits === 0 ? 0 : (0xffffffff << (32 - bits)) >>> 0;
    return (n & mask) === (b & mask);
  };
  return (
    inRange("0.0.0.0", 8) || // "this" network
    inRange("10.0.0.0", 8) || // private
    inRange("100.64.0.0", 10) || // CGNAT
    inRange("127.0.0.0", 8) || // loopback
    inRange("169.254.0.0", 16) || // link-local (incl. cloud metadata)
    inRange("172.16.0.0", 12) || // private
    inRange("192.0.0.0", 24) || // IETF protocol assignments
    inRange("192.168.0.0", 16) || // private
    inRange("198.18.0.0", 15) || // benchmarking
    inRange("224.0.0.0", 4) || // multicast
    inRange("240.0.0.0", 4) // reserved / broadcast
  );
}

/** Expand an IPv6 string to its 16 bytes, handling `::` and embedded IPv4. */
function ipv6ToBytes(ip: string): number[] | undefined {
  let s = ip;
  // Strip zone id (fe80::1%eth0) and brackets.
  s = s.replace(/^\[|\]$/g, "").replace(/%.*$/, "");
  const halves = s.split("::");
  if (halves.length > 2) return undefined;

  const parseGroups = (segment: string): number[] | undefined => {
    if (segment === "") return [];
    const out: number[] = [];
    for (const part of segment.split(":")) {
      if (part.includes(".")) {
        // Embedded IPv4 tail.
        const v4 = ipv4ToInt(part);
        if (v4 === undefined) return undefined;
        out.push((v4 >>> 16) & 0xffff, v4 & 0xffff);
      } else {
        if (!/^[0-9a-fA-F]{1,4}$/.test(part)) return undefined;
        out.push(parseInt(part, 16));
      }
    }
    return out;
  };

  const head = parseGroups(halves[0]);
  const tail = halves.length === 2 ? parseGroups(halves[1]) : [];
  if (!head || !tail) return undefined;

  let groups: number[];
  if (halves.length === 2) {
    const fill = 8 - head.length - tail.length;
    if (fill < 0) return undefined;
    groups = [...head, ...new Array(fill).fill(0), ...tail];
  } else {
    groups = head;
  }
  if (groups.length !== 8) return undefined;
  const bytes: number[] = [];
  for (const g of groups) bytes.push((g >>> 8) & 0xff, g & 0xff);
  return bytes;
}

function isBlockedIPv6(ip: string): boolean {
  const b = ipv6ToBytes(ip);
  if (!b) return true; // unparseable -> unsafe

  // IPv4-mapped (::ffff:a.b.c.d) and IPv4-compatible: classify as IPv4.
  const firstTenZero = b.slice(0, 10).every((x) => x === 0);
  if (firstTenZero && b[10] === 0xff && b[11] === 0xff) {
    return isBlockedIPv4(`${b[12]}.${b[13]}.${b[14]}.${b[15]}`);
  }

  const allZero = b.every((x) => x === 0);
  if (allZero) return true; // :: unspecified
  if (b.slice(0, 15).every((x) => x === 0) && b[15] === 1) return true; // ::1 loopback
  if ((b[0] & 0xfe) === 0xfc) return true; // fc00::/7 unique-local
  if (b[0] === 0xfe && (b[1] & 0xc0) === 0x80) return true; // fe80::/10 link-local
  if (b[0] === 0xff) return true; // ff00::/8 multicast
  return false;
}

/** True when `ip` (v4 or v6 literal) is loopback/private/link-local/etc. */
export function isBlockedAddress(ip: string): boolean {
  const kind = net.isIP(ip);
  if (kind === 4) return isBlockedIPv4(ip);
  if (kind === 6) return isBlockedIPv6(ip);
  return true; // not a valid IP literal -> unsafe
}

// ---------------------------------------------------------------------------
// URL / host validation
// ---------------------------------------------------------------------------

export class BlockedUrlError extends Error {}

/**
 * Validate a single URL: http/https scheme, and (unless private networking is
 * allowed) every DNS-resolved address must be public. Throws BlockedUrlError on
 * rejection. Returns the parsed URL on success.
 */
export async function assertSafeUrl(rawUrl: string, allowPrivateNetwork: boolean): Promise<URL> {
  let u: URL;
  try {
    u = new URL(rawUrl);
  } catch {
    throw new BlockedUrlError(`invalid URL: ${rawUrl}`);
  }
  if (u.protocol !== "http:" && u.protocol !== "https:") {
    throw new BlockedUrlError(`unsupported scheme "${u.protocol}" (only http/https allowed)`);
  }
  if (allowPrivateNetwork) return u;

  const host = u.hostname.replace(/^\[|\]$/g, "");
  // If the host is already an IP literal, check it directly.
  if (net.isIP(host)) {
    if (isBlockedAddress(host)) {
      throw new BlockedUrlError(`blocked address ${host} (loopback/private/link-local)`);
    }
    return u;
  }
  // Otherwise resolve and reject if ANY resolved address is private.
  let addrs: { address: string }[];
  try {
    addrs = await dns.lookup(host, { all: true });
  } catch {
    throw new BlockedUrlError(`could not resolve host ${host}`);
  }
  if (addrs.length === 0) throw new BlockedUrlError(`host ${host} did not resolve`);
  for (const a of addrs) {
    if (isBlockedAddress(a.address)) {
      throw new BlockedUrlError(`host ${host} resolves to blocked address ${a.address}`);
    }
  }
  return u;
}

/**
 * Boolean SSRF predicate: true when `host` (an IP literal or hostname) should be
 * refused. Shares the same `isBlockedAddress` classifier as `assertSafeUrl`, so
 * the browser-render path's per-request route handler (which sees raw request
 * hosts, not URLs) classifies identically. A resolution failure is treated as
 * blocked — the safe default for an interceptor that must decide synchronously.
 */
export async function hostResolvesToBlocked(host: string, allowPrivate: boolean): Promise<boolean> {
  if (allowPrivate) return false;
  const h = host.replace(/^\[|\]$/g, "");
  if (net.isIP(h)) return isBlockedAddress(h);
  let addrs: { address: string }[];
  try {
    addrs = await dns.lookup(h, { all: true });
  } catch {
    return true;
  }
  if (addrs.length === 0) return true;
  return addrs.some((a) => isBlockedAddress(a.address));
}

// ---------------------------------------------------------------------------
// Guarded fetch with manual redirect re-validation
// ---------------------------------------------------------------------------

export interface SafeFetchOptions {
  headers?: Record<string, string>;
  method?: string;
  body?: string;
  timeoutMs: number;
  signal?: AbortSignal;
  allowPrivateNetwork?: boolean;
  maxRedirects?: number;
}

const DEFAULT_MAX_REDIRECTS = 5;

/**
 * fetch() with a per-request timeout and SSRF guard. Redirects are followed
 * manually so each hop's target is re-validated before it is requested.
 */
export async function safeFetch(rawUrl: string, opts: SafeFetchOptions): Promise<Response> {
  const allowPrivate = opts.allowPrivateNetwork === true;
  const maxRedirects = opts.maxRedirects ?? DEFAULT_MAX_REDIRECTS;
  let currentUrl = rawUrl;

  for (let hop = 0; hop <= maxRedirects; hop++) {
    const u = await assertSafeUrl(currentUrl, allowPrivate);

    const ctrl = new AbortController();
    const timer = setTimeout(() => ctrl.abort(), opts.timeoutMs);
    const onAbort = () => ctrl.abort();
    opts.signal?.addEventListener("abort", onAbort, { once: true });
    let res: Response;
    try {
      res = await fetch(u, {
        method: opts.method,
        body: opts.body,
        headers: opts.headers,
        redirect: "manual",
        signal: ctrl.signal,
      });
    } finally {
      clearTimeout(timer);
      opts.signal?.removeEventListener("abort", onAbort);
    }

    // Manual redirect: re-validate the next hop before following it.
    if (res.status >= 300 && res.status < 400) {
      const location = res.headers.get("location");
      if (!location) return res; // 3xx with no Location: hand back as-is
      if (hop === maxRedirects) {
        throw new BlockedUrlError(`too many redirects (>${maxRedirects})`);
      }
      currentUrl = new URL(location, u).toString();
      // Drain the redirect body so the socket can be reused.
      await res.body?.cancel().catch(() => {});
      continue;
    }
    return res;
  }
  // Unreachable, but keeps the type checker happy.
  throw new BlockedUrlError("redirect handling exhausted");
}

// ---------------------------------------------------------------------------
// Bounded body reads
// ---------------------------------------------------------------------------

export class ResponseTooLargeError extends Error {}

/**
 * Read a response body as text, aborting once `maxBytes` is exceeded. Honors an
 * oversized Content-Length up front. Decodes as UTF-8.
 */
export async function readCappedText(res: Response, maxBytes: number): Promise<string> {
  const declared = Number(res.headers.get("content-length"));
  if (Number.isFinite(declared) && declared > maxBytes) {
    await res.body?.cancel().catch(() => {});
    throw new ResponseTooLargeError(`response too large (${declared} bytes > ${maxBytes} cap)`);
  }
  if (!res.body) return "";

  const reader = res.body.getReader();
  const decoder = new TextDecoder("utf-8");
  let received = 0;
  let out = "";
  try {
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      received += value.byteLength;
      if (received > maxBytes) {
        throw new ResponseTooLargeError(`response exceeded ${maxBytes} byte cap`);
      }
      out += decoder.decode(value, { stream: true });
    }
    out += decoder.decode();
    return out;
  } finally {
    reader.cancel().catch(() => {});
  }
}

// ---------------------------------------------------------------------------
// High-level fetchers used by the tools
// ---------------------------------------------------------------------------

// Body-size caps (post-decode is bounded by these pre-decode byte caps).
export const MAX_HTML_BYTES = 3_000_000;
export const MAX_TEXT_FILE_BYTES = 262_144; // 256 KB for llms.txt probes
export const MAX_SNIFF_BYTES = 65_536; // 64 KB read from a failed body to detect a bot-challenge

// ---------------------------------------------------------------------------
// Browser-like headers (Tier 0 bot-evasion)
// ---------------------------------------------------------------------------

const DEFAULT_ACCEPT_LANGUAGE = "en-US,en;q=0.9";

/** A real browser sends client-hints only from Chromium; the default UA is Firefox. */
function isChromiumUA(ua: string): boolean {
  return /(Chrome|Chromium|Edg)\//.test(ua) && !/Firefox\//.test(ua);
}

function originRoot(url: string): string | undefined {
  try {
    return `${new URL(url).origin}/`;
  } catch {
    return undefined;
  }
}

/**
 * Build realistic browser request headers. Sending a fuller, consistent header
 * set (Accept-Language, Sec-Fetch-*, Upgrade-Insecure-Requests, Referer) gets
 * past naive UA-only bot checks before we pay for a headless browser. Chromium
 * `Sec-Ch-Ua*` client-hints are emitted ONLY for a Chromium-looking UA — pairing
 * them with the Firefox default would itself be a bot tell.
 */
export function browserHeaders(
  userAgent: string,
  acceptLanguage: string = DEFAULT_ACCEPT_LANGUAGE,
  opts?: { accept?: string; referer?: string },
): Record<string, string> {
  const h: Record<string, string> = {
    "User-Agent": userAgent,
    Accept: opts?.accept ?? "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
    "Accept-Language": acceptLanguage || DEFAULT_ACCEPT_LANGUAGE,
    "Sec-Fetch-Dest": "document",
    "Sec-Fetch-Mode": "navigate",
    "Sec-Fetch-Site": "none",
    "Sec-Fetch-User": "?1",
    "Upgrade-Insecure-Requests": "1",
  };
  if (opts?.referer) h.Referer = opts.referer;
  if (isChromiumUA(userAgent)) {
    h["Sec-Ch-Ua"] = '"Chromium";v="128", "Not;A=Brand";v="24", "Google Chrome";v="128"';
    h["Sec-Ch-Ua-Mobile"] = "?0";
    h["Sec-Ch-Ua-Platform"] = '"Linux"';
  }
  return h;
}

// ---------------------------------------------------------------------------
// Bot-block / challenge classification (pure)
// ---------------------------------------------------------------------------

export interface BlockClassification {
  blocked: boolean;
  reason?: string; // cloudflare-challenge | rate-limited | forbidden-bot | perimeterx | datadome | ...
}

/**
 * Decide whether a response is a bot-block / interstitial (vs. a plain failure).
 * Pure and unit-testable. Conservative on HTTP 200: only strong script/header
 * markers (which appear only on genuine challenge pages) count, so an ordinary
 * page that merely mentions Cloudflare isn't mis-flagged. Softer text phrases are
 * trusted only alongside a block-ish status (403/429/503). A plain 401 (auth),
 * 404, or 5xx without a challenge marker is NOT a block — keeping the block
 * ledger free of false positives that would warn the agent off a good domain.
 */
export function classifyBlock(
  status: number,
  headers: Headers | Record<string, string>,
  body: string,
): BlockClassification {
  const get = (k: string): string => {
    if (headers instanceof Headers) return headers.get(k) ?? "";
    return headers[k] ?? headers[k.toLowerCase()] ?? "";
  };
  const cfMitigated = get("cf-mitigated").toLowerCase();
  const cfRay = get("cf-ray");
  const server = get("server").toLowerCase();
  const isCf = server.includes("cloudflare") || cfRay !== "";
  const b = body.toLowerCase();

  // Strong markers — reliable even on HTTP 200 (present only on challenge pages).
  if (cfMitigated === "challenge") return { blocked: true, reason: "cloudflare-challenge" };
  if (b.includes("_cf_chl_opt") || b.includes("cf-browser-verification") || b.includes("/cdn-cgi/challenge-platform")) {
    return { blocked: true, reason: "cloudflare-challenge" };
  }
  if (b.includes("px-captcha") || b.includes("_pxhd")) return { blocked: true, reason: "perimeterx" };
  if (b.includes("datadome") && (b.includes("captcha") || status === 403)) return { blocked: true, reason: "datadome" };

  const candidate = status === 403 || status === 429 || status === 503;
  if (!candidate) return { blocked: false };

  // Softer text signatures: trusted only alongside a block-ish status.
  if (
    b.includes("just a moment") ||
    b.includes("enable javascript and cookies to continue") ||
    b.includes("verify you are a human") ||
    b.includes("are you a robot") ||
    (b.includes("attention required") && isCf)
  ) {
    return { blocked: true, reason: "cloudflare-challenge" };
  }
  if (status === 429) return { blocked: true, reason: "rate-limited" };
  if (status === 403) return { blocked: true, reason: "forbidden-bot" };
  if (status === 503 && isCf) return { blocked: true, reason: "cloudflare-unavailable" };
  return { blocked: false };
}

export interface FetchHtmlResult {
  html?: string;
  error?: string;
  blocked?: true;
  status?: number;
  blockReason?: string;
}

export interface FetchHtmlOptions {
  timeoutMs: number;
  userAgent: string;
  acceptLanguage?: string;
  allowPrivateNetwork?: boolean;
  signal?: AbortSignal;
}

/**
 * Fetch a page as HTML text (guarded + size-capped). A response detected as a
 * bot-block/challenge (see `classifyBlock`) returns `{ blocked: true, ... }` so
 * callers can escalate (headless browser / reader) rather than give up. Other
 * failures return a plain `{ error }` and must NOT be treated as a domain block.
 */
export async function fetchHtml(url: string, opts: FetchHtmlOptions): Promise<FetchHtmlResult> {
  try {
    const res = await safeFetch(url, {
      headers: browserHeaders(opts.userAgent, opts.acceptLanguage, { referer: originRoot(url) }),
      timeoutMs: opts.timeoutMs,
      signal: opts.signal,
      allowPrivateNetwork: opts.allowPrivateNetwork,
    });
    if (!res.ok) {
      const candidate = res.status === 403 || res.status === 429 || res.status === 503;
      let bodySlice = "";
      if (candidate) {
        try {
          bodySlice = await readCappedText(res, MAX_SNIFF_BYTES);
        } catch {
          /* body unreadable — classify on status/headers alone */
        }
      } else {
        await res.body?.cancel().catch(() => {});
      }
      const cls = classifyBlock(res.status, res.headers, bodySlice);
      if (cls.blocked) {
        return {
          blocked: true,
          status: res.status,
          blockReason: cls.reason,
          error: `blocked: ${cls.reason ?? "bot protection"} (HTTP ${res.status})`,
        };
      }
      return { error: `HTTP ${res.status} ${res.statusText}`, status: res.status };
    }
    const ct = res.headers.get("content-type") ?? "";
    if (!/text\/html|application\/xhtml/i.test(ct)) {
      await res.body?.cancel().catch(() => {});
      return { error: `unsupported content-type: ${ct || "unknown"}` };
    }
    const html = await readCappedText(res, MAX_HTML_BYTES);
    // Cloudflare sometimes serves the interstitial with HTTP 200 — re-check.
    const cls = classifyBlock(res.status, res.headers, html.slice(0, MAX_SNIFF_BYTES));
    if (cls.blocked) {
      return { blocked: true, status: res.status, blockReason: cls.reason, error: `blocked: ${cls.reason ?? "bot protection"} (HTTP ${res.status})` };
    }
    return { html };
  } catch (e) {
    return { error: classifyFetchError(e) };
  }
}

export interface FetchTextFileOptions {
  timeoutMs: number;
  userAgent: string;
  acceptLanguage?: string;
  allowPrivateNetwork?: boolean;
  signal?: AbortSignal;
}

/**
 * Fetch a candidate text file (e.g. llms.txt), gated by content-type. Returns
 * the raw text, or undefined when absent/rejected. Throws "aborted" only when
 * the caller's signal fired (so callers can avoid negative-caching a cancel).
 */
export async function fetchTextFile(url: string, opts: FetchTextFileOptions): Promise<string | undefined> {
  try {
    const res = await safeFetch(url, {
      headers: browserHeaders(opts.userAgent, opts.acceptLanguage, { accept: "text/markdown,text/plain,*/*" }),
      timeoutMs: opts.timeoutMs,
      signal: opts.signal,
      allowPrivateNetwork: opts.allowPrivateNetwork,
    });
    if (!res.ok) {
      await res.body?.cancel().catch(() => {});
      return undefined;
    }
    const ct = (res.headers.get("content-type") ?? "").toLowerCase();
    // Reject obvious non-text types. Crucially reject text/html so a 200
    // soft-404 page is never injected into context as if it were llms.txt.
    if (/text\/html|application\/(json|xml|octet-stream)|image\//.test(ct)) {
      await res.body?.cancel().catch(() => {});
      return undefined;
    }
    return (await readCappedText(res, MAX_TEXT_FILE_BYTES)).trim();
  } catch (e) {
    if (opts.signal?.aborted) throw new Error("aborted");
    if (e instanceof ResponseTooLargeError) return undefined; // oversized llms.txt: skip quietly
    return undefined;
  }
}

function classifyFetchError(e: unknown): string {
  if (e instanceof BlockedUrlError) return `blocked: ${e.message}`;
  if (e instanceof ResponseTooLargeError) return e.message;
  const err = e as { name?: string; message?: string };
  if (err?.name === "AbortError") return "timeout/aborted";
  return err?.message || String(e);
}