magi-code 0.77.1

use super::contract::metadata_key as meta;
use crate::{cancellation::AgentCancellation, output::redact_sensitive_text};
use dom_smoothie::{Config, Readability, TextMode};
use encoding_rs::{Encoding, UTF_8};
use reqwest::{StatusCode, Url, blocking::Client, redirect::Policy};
use serde_json::json;
use std::{
    io::Read,
    net::{IpAddr, SocketAddr, ToSocketAddrs},
    sync::{
        atomic::{AtomicUsize, Ordering},
        mpsc,
    },
    thread::{self, JoinHandle},
    time::{Duration, Instant},
};

const URL_FETCH_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
const URL_FETCH_DNS_TIMEOUT: Duration = Duration::from_secs(5);
const URL_FETCH_MAX_DNS_RESOLVER_WORKERS: usize = 8;
const URL_FETCH_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
const URL_FETCH_MAX_BYTES: usize = 512 * 1024;
const URL_FETCH_OUTPUT_MAX_BYTES: usize = 48 * 1024;
const URL_FETCH_DEFAULT_MAX_TOKENS: u64 = 5_000;
const URL_FETCH_MIN_MAX_TOKENS: u64 = 1_000;
const URL_FETCH_MAX_MAX_TOKENS: u64 = 10_000;
const URL_FETCH_READABILITY_MIN_LENGTH: usize = 200;
const TRUNCATION_MARKER: &str = "\n[truncated]";
const URL_FETCH_CANCEL_POLL_INTERVAL: Duration = Duration::from_millis(25);
const URL_FETCH_WORKER_JOIN_TIMEOUT: Duration = Duration::from_millis(500);

static URL_FETCH_DNS_RESOLVER_WORKER_COUNT: AtomicUsize = AtomicUsize::new(0);

#[derive(Debug, Clone)]
struct PinnedTarget {
    host: String,
    addrs: Vec<SocketAddr>,
}

#[derive(Debug, Clone)]
struct ExtractedContent {
    text: String,
    title: Option<String>,
}

#[derive(Debug, Clone)]
pub(crate) struct UrlFetchInput {
    pub(crate) url: String,
    pub(crate) max_tokens: Option<u64>,
}

impl UrlFetchInput {
    pub(crate) fn validate(mut self) -> anyhow::Result<Self> {
        self.url = self.url.trim().to_string();
        if self.url.is_empty() {
            anyhow::bail!("url must not be empty");
        }
        let parsed = Url::parse(&self.url)
            .map_err(|_| anyhow::anyhow!("url must be an http(s) URL with a host"))?;
        if !matches!(parsed.scheme(), "http" | "https") || parsed.host_str().is_none() {
            anyhow::bail!("url must be an http(s) URL with a host");
        }
        if let Some(max_tokens) = self.max_tokens
            && !(URL_FETCH_MIN_MAX_TOKENS..=URL_FETCH_MAX_MAX_TOKENS).contains(&max_tokens)
        {
            anyhow::bail!(
                "maxTokens must be between {URL_FETCH_MIN_MAX_TOKENS} and {URL_FETCH_MAX_MAX_TOKENS}"
            );
        }
        Ok(self)
    }

    fn max_tokens(&self) -> u64 {
        self.max_tokens.unwrap_or(URL_FETCH_DEFAULT_MAX_TOKENS)
    }
}

#[derive(Debug, Clone)]
pub(crate) struct UrlFetchOutput {
    pub(crate) success: bool,
    pub(crate) content: String,
    pub(crate) metadata: serde_json::Value,
}

pub(crate) fn fetch_url(
    input: UrlFetchInput,
    cancellation: &AgentCancellation,
) -> anyhow::Result<UrlFetchOutput> {
    cancellation.check()?;
    let url = Url::parse(&input.url)?;
    let response = fetch_protected_response(url.clone(), cancellation)?;
    fetch_with_response(input, url, response, cancellation)
}

fn build_pinned_client(target: &PinnedTarget) -> anyhow::Result<Client> {
    Client::builder()
        .connect_timeout(URL_FETCH_CONNECT_TIMEOUT)
        .timeout(URL_FETCH_REQUEST_TIMEOUT)
        .user_agent(format!("magi-code/{}", env!("CARGO_PKG_VERSION")))
        .redirect(Policy::none())
        .no_proxy()
        .resolve_to_addrs(&target.host, &target.addrs)
        .build()
        .map_err(|error| anyhow::anyhow!("url_fetch HTTP client setup failed: {error}"))
}

fn fetch_protected_response(
    url: Url,
    cancellation: &AgentCancellation,
) -> anyhow::Result<FetchHttpResponse> {
    cancellation.check()?;
    let target = resolve_public_target(&url, cancellation)?;
    cancellation.check()?;
    let client = build_pinned_client(&target)?;
    cancellation.check()?;
    fetch_response_with_client(&client, &url, cancellation)
}

#[derive(Debug)]
struct FetchHttpResponse {
    status: StatusCode,
    content_type: String,
    bytes: Vec<u8>,
}

fn is_html_content_type(content_type: &str) -> bool {
    let media_type = content_type.split(';').next().map(str::trim);
    media_type.is_some_and(|media_type| {
        media_type.eq_ignore_ascii_case("text/html")
            || media_type.eq_ignore_ascii_case("application/xhtml+xml")
    })
}

fn decode_response_body(bytes: &[u8], content_type: &str) -> String {
    if let Some((encoding, _)) = Encoding::for_bom(bytes) {
        return encoding.decode(bytes).0.into_owned();
    }

    if let Some(encoding) =
        declared_charset(content_type).and_then(|label| Encoding::for_label(label.as_bytes()))
    {
        return encoding.decode_without_bom_handling(bytes).0.into_owned();
    }

    if is_html_content_type(content_type)
        && let Some(label) = sniff_html_meta_charset(bytes)
        && let Some(encoding) = meta_charset_encoding(&label)
    {
        return encoding.decode_without_bom_handling(bytes).0.into_owned();
    }

    UTF_8.decode_without_bom_handling(bytes).0.into_owned()
}

fn declared_charset(content_type: &str) -> Option<String> {
    content_type.split(';').skip(1).find_map(|parameter| {
        let (name, value) = parameter.split_once('=')?;
        if !name.trim().eq_ignore_ascii_case("charset") {
            return None;
        }
        let value = value.trim().trim_matches(['"', '\'']).trim();
        (!value.is_empty()).then(|| value.to_string())
    })
}

fn sniff_html_meta_charset(bytes: &[u8]) -> Option<String> {
    let prefix = &bytes[..bytes.len().min(1024)];
    let lower = prefix
        .iter()
        .map(|byte| byte.to_ascii_lowercase())
        .collect::<Vec<_>>();
    let mut cursor = 0;
    while cursor + b"<meta".len() <= lower.len() {
        if lower[cursor..].starts_with(b"<meta")
            && lower
                .get(cursor + b"<meta".len())
                .is_none_or(|byte| byte.is_ascii_whitespace() || *byte == b'/' || *byte == b'>')
        {
            let attributes_start = cursor + b"<meta".len();
            let tag_end = lower[attributes_start..]
                .iter()
                .position(|byte| *byte == b'>')
                .map_or(lower.len(), |offset| attributes_start + offset);
            if let Some(label) = charset_value(&lower[attributes_start..tag_end]) {
                return Some(label);
            }
            cursor = tag_end.saturating_add(1);
        } else {
            cursor += 1;
        }
    }
    None
}

fn charset_value(attributes: &[u8]) -> Option<String> {
    let mut cursor = 0;
    while cursor + b"charset".len() <= attributes.len() {
        if attributes[cursor..].starts_with(b"charset")
            && (cursor == 0 || !is_attribute_name_byte(attributes[cursor - 1]))
            && (cursor + b"charset".len() == attributes.len()
                || !is_attribute_name_byte(attributes[cursor + b"charset".len()]))
        {
            let mut value_start = cursor + b"charset".len();
            while attributes
                .get(value_start)
                .is_some_and(u8::is_ascii_whitespace)
            {
                value_start += 1;
            }
            if attributes.get(value_start) == Some(&b'=') {
                value_start += 1;
                while attributes
                    .get(value_start)
                    .is_some_and(u8::is_ascii_whitespace)
                {
                    value_start += 1;
                }
                let quote = attributes.get(value_start).copied();
                let (value_start, value_end) = match quote {
                    Some(quote @ (b'"' | b'\'')) => {
                        let value_start = value_start + 1;
                        let value_end = attributes[value_start..]
                            .iter()
                            .position(|byte| *byte == quote)
                            .map_or(attributes.len(), |offset| value_start + offset);
                        (value_start, value_end)
                    }
                    _ => {
                        let value_end = attributes[value_start..]
                            .iter()
                            .position(|byte| {
                                byte.is_ascii_whitespace() || matches!(*byte, b';' | b'/')
                            })
                            .map_or(attributes.len(), |offset| value_start + offset);
                        (value_start, value_end)
                    }
                };
                if value_start < value_end {
                    return Some(
                        String::from_utf8_lossy(&attributes[value_start..value_end]).into_owned(),
                    );
                }
            }
        }
        cursor += 1;
    }
    None
}

fn is_attribute_name_byte(byte: u8) -> bool {
    byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b':')
}

fn meta_charset_encoding(label: &str) -> Option<&'static Encoding> {
    if label.eq_ignore_ascii_case("utf-16")
        || label.eq_ignore_ascii_case("utf-16le")
        || label.eq_ignore_ascii_case("utf-16be")
    {
        // HTML's prescan treats UTF-16 meta declarations as UTF-8. A UTF-16
        // document with a BOM is handled above, before this meta fallback.
        return Some(UTF_8);
    }
    Encoding::for_label(label.as_bytes())
}

struct DnsResolverPermit<'a> {
    counter: &'a AtomicUsize,
}

impl<'a> DnsResolverPermit<'a> {
    fn try_acquire(counter: &'a AtomicUsize, limit: usize) -> Option<Self> {
        if limit == 0 {
            return None;
        }

        let mut current = counter.load(Ordering::Acquire);
        loop {
            if current >= limit {
                return None;
            }
            match counter.compare_exchange_weak(
                current,
                current + 1,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => return Some(Self { counter }),
                Err(observed) => current = observed,
            }
        }
    }
}

impl Drop for DnsResolverPermit<'_> {
    fn drop(&mut self) {
        self.counter.fetch_sub(1, Ordering::Release);
    }
}

fn acquire_dns_resolver_permit<'a>(
    counter: &'a AtomicUsize,
    limit: usize,
) -> anyhow::Result<DnsResolverPermit<'a>> {
    DnsResolverPermit::try_acquire(counter, limit)
        .ok_or_else(|| anyhow::anyhow!("url_fetch DNS resolver capacity exhausted; retry later"))
}

struct UrlFetchWorkerHandle {
    worker_label: &'static str,
    done_receiver: mpsc::Receiver<()>,
    join_handle: JoinHandle<()>,
}

impl UrlFetchWorkerHandle {
    fn join_or_warn(self) {
        match self
            .done_receiver
            .recv_timeout(URL_FETCH_WORKER_JOIN_TIMEOUT)
        {
            Ok(()) | Err(mpsc::RecvTimeoutError::Disconnected) => {
                if self.join_handle.join().is_err() {
                    eprintln!(
                        "magi-code warning: url_fetch {} worker panicked",
                        self.worker_label
                    );
                }
            }
            Err(mpsc::RecvTimeoutError::Timeout) => {
                eprintln!(
                    "magi-code warning: url_fetch {} worker did not exit within cleanup grace period; detaching worker",
                    self.worker_label
                );
            }
        }
    }

    fn detach_with_warning(self) {
        let detail = match self.worker_label {
            "DNS" => {
                "worker detached; the OS resolver may continue running after the caller deadline"
            }
            _ => "worker detached; it may continue until its own timeout",
        };
        eprintln!(
            "magi-code warning: url_fetch {} {}",
            self.worker_label, detail
        );
        let UrlFetchWorkerHandle { join_handle, .. } = self;
        drop(join_handle);
    }
}

fn spawn_url_fetch_worker(
    client: Client,
    url: Url,
) -> anyhow::Result<(
    mpsc::Receiver<anyhow::Result<FetchHttpResponse>>,
    UrlFetchWorkerHandle,
)> {
    let (sender, receiver) = mpsc::sync_channel(1);
    let (done_sender, done_receiver) = mpsc::sync_channel(1);
    let join_handle = thread::Builder::new()
        .name("url-fetch-http".to_string())
        .spawn(move || {
            let _ = sender.send(blocking_fetch_response(client, url));
            let _ = done_sender.send(());
        })?;
    Ok((
        receiver,
        UrlFetchWorkerHandle {
            worker_label: "HTTP",
            done_receiver,
            join_handle,
        },
    ))
}

fn spawn_dns_resolution_worker(
    url: Url,
) -> anyhow::Result<(
    mpsc::Receiver<anyhow::Result<PinnedTarget>>,
    UrlFetchWorkerHandle,
)> {
    let permit = acquire_dns_resolver_permit(
        &URL_FETCH_DNS_RESOLVER_WORKER_COUNT,
        URL_FETCH_MAX_DNS_RESOLVER_WORKERS,
    )?;
    let (sender, receiver) = mpsc::sync_channel(1);
    let (done_sender, done_receiver) = mpsc::sync_channel(1);
    let join_handle = thread::Builder::new()
        .name("url-fetch-dns".to_string())
        .spawn(move || {
            let _permit = permit;
            let _ = sender.send(resolve_public_target_sync(&url));
            let _ = done_sender.send(());
        })?;
    Ok((
        receiver,
        UrlFetchWorkerHandle {
            worker_label: "DNS",
            done_receiver,
            join_handle,
        },
    ))
}

fn recv_cancellable<T>(
    receiver: &mpsc::Receiver<T>,
    timeout_label: &str,
    idle_timeout: Duration,
    cancellation: &AgentCancellation,
) -> anyhow::Result<T> {
    let start = Instant::now();
    loop {
        cancellation.check()?;
        if start.elapsed() >= idle_timeout {
            anyhow::bail!("{timeout_label} after {}s", idle_timeout.as_secs());
        }
        let remaining = idle_timeout.saturating_sub(start.elapsed());
        match receiver.recv_timeout(remaining.min(URL_FETCH_CANCEL_POLL_INTERVAL)) {
            Ok(value) => return Ok(value),
            Err(mpsc::RecvTimeoutError::Timeout) => continue,
            Err(mpsc::RecvTimeoutError::Disconnected) => {
                cancellation.check()?;
                anyhow::bail!("url_fetch worker disconnected");
            }
        }
    }
}

fn wait_for_worker_result<T>(
    receiver: mpsc::Receiver<anyhow::Result<T>>,
    worker: UrlFetchWorkerHandle,
    timeout_label: &str,
    timeout: Duration,
    cancellation: &AgentCancellation,
) -> anyhow::Result<T> {
    let result = match recv_cancellable(&receiver, timeout_label, timeout, cancellation) {
        Ok(result) => result,
        Err(error) => {
            // A blocking resolver or reqwest call cannot be force-aborted here. Drop the
            // bounded result receiver and detach immediately; do not add cleanup grace after
            // the caller deadline. The DNS worker cap bounds resolver threads even when the OS
            // resolver outlives the caller deadline; its permit remains held until worker exit.
            drop(receiver);
            worker.detach_with_warning();
            return Err(error);
        }
    };
    worker.join_or_warn();
    result
}

fn blocking_fetch_response(client: Client, url: Url) -> anyhow::Result<FetchHttpResponse> {
    let response = client.get(url).send().map_err(|error| {
        if error.is_timeout() {
            anyhow::anyhow!("url_fetch request timed out")
        } else {
            anyhow::anyhow!("url_fetch network error: {error}")
        }
    })?;
    let status = response.status();
    let content_type = response
        .headers()
        .get(reqwest::header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .unwrap_or_default()
        .to_string();
    if response
        .content_length()
        .is_some_and(|length| length > URL_FETCH_MAX_BYTES as u64)
    {
        anyhow::bail!("url_fetch response exceeded {URL_FETCH_MAX_BYTES} byte limit");
    }
    let bytes = read_bounded_body(response)?;
    Ok(FetchHttpResponse {
        status,
        content_type,
        bytes,
    })
}

fn fetch_response_with_client(
    client: &Client,
    url: &Url,
    cancellation: &AgentCancellation,
) -> anyhow::Result<FetchHttpResponse> {
    let (receiver, worker) = spawn_url_fetch_worker(client.clone(), url.clone())?;
    let response = wait_for_worker_result(
        receiver,
        worker,
        "url_fetch response timeout",
        URL_FETCH_REQUEST_TIMEOUT,
        cancellation,
    )?;
    cancellation.check()?;
    Ok(response)
}

#[cfg(test)]
fn fetch_with_client(
    args: UrlFetchInput,
    url: Url,
    client: &Client,
    cancellation: &AgentCancellation,
) -> anyhow::Result<UrlFetchOutput> {
    let response = fetch_response_with_client(client, &url, cancellation)?;
    fetch_with_response(args, url, response, cancellation)
}

fn fetch_with_response(
    args: UrlFetchInput,
    url: Url,
    response: FetchHttpResponse,
    cancellation: &AgentCancellation,
) -> anyhow::Result<UrlFetchOutput> {
    let status = response.status;
    let content_type = response.content_type;
    let byte_count = response.bytes.len();
    let raw = decode_response_body(&response.bytes, &content_type);
    if !status.is_success() {
        let (content, truncated) = sanitize_and_bound(&raw, args.max_tokens());
        return Ok(result(FetchResultParts {
            success: false,
            content,
            url: &url,
            status,
            content_type: &content_type,
            title: None,
            bytes: byte_count,
            truncated,
        }));
    }
    let extracted = extract_content(&raw, &content_type, url.as_str());
    cancellation.check()?;
    let (content, truncated) = sanitize_and_bound(&extracted.text, args.max_tokens());
    cancellation.check()?;
    Ok(result(FetchResultParts {
        success: true,
        content,
        url: &url,
        status,
        content_type: &content_type,
        title: extracted.title.as_deref(),
        bytes: byte_count,
        truncated,
    }))
}

struct FetchResultParts<'a> {
    success: bool,
    content: String,
    url: &'a Url,
    status: StatusCode,
    content_type: &'a str,
    title: Option<&'a str>,
    bytes: usize,
    truncated: bool,
}

fn result(parts: FetchResultParts<'_>) -> UrlFetchOutput {
    UrlFetchOutput {
        success: parts.success,
        content: parts.content,
        metadata: json!({
            (meta::URL): parts.url.as_str(),
            (meta::STATUS): parts.status.as_u16(),
            (meta::CONTENT_TYPE): parts.content_type,
            (meta::TITLE): parts.title,
            (meta::BYTES): parts.bytes,
            (meta::TRUNCATED): parts.truncated,
        }),
    }
}

fn read_bounded_body(mut response: reqwest::blocking::Response) -> anyhow::Result<Vec<u8>> {
    let mut bytes = Vec::new();
    response
        .by_ref()
        .take((URL_FETCH_MAX_BYTES + 1) as u64)
        .read_to_end(&mut bytes)
        .map_err(|error| anyhow::anyhow!("url_fetch response read failed: {error}"))?;
    if bytes.len() > URL_FETCH_MAX_BYTES {
        anyhow::bail!("url_fetch response exceeded {URL_FETCH_MAX_BYTES} byte limit");
    }
    Ok(bytes)
}

fn extract_content(body: &str, content_type: &str, url: &str) -> ExtractedContent {
    if !content_type.to_ascii_lowercase().contains("html") {
        return ExtractedContent {
            text: body.to_string(),
            title: None,
        };
    }
    if let Some(extracted) = readability_extract(body, url) {
        return extracted;
    }
    ExtractedContent {
        text: html2md::rewrite_html(body, false),
        title: None,
    }
}

fn readability_extract(html: &str, url: &str) -> Option<ExtractedContent> {
    let cfg = Config {
        text_mode: TextMode::Markdown,
        ..Default::default()
    };
    let mut readability = Readability::new(html, Some(url), Some(cfg)).ok()?;
    let article = readability.parse().ok()?;
    let text = article.text_content.to_string();
    let text = text.trim().to_string();
    if article.length <= URL_FETCH_READABILITY_MIN_LENGTH || text.is_empty() {
        return None;
    }
    let title = (!article.title.trim().is_empty()).then(|| article.title.trim().to_string());
    Some(ExtractedContent { text, title })
}

fn sanitize_and_bound(text: &str, max_tokens: u64) -> (String, bool) {
    let cap = URL_FETCH_OUTPUT_MAX_BYTES.min(max_tokens.saturating_mul(4) as usize);
    let mut truncated = false;
    let mut content = truncate_with_marker(text.trim(), cap, &mut truncated);
    content = redact_sensitive_text(&content);
    if content.len() > cap {
        truncated = true;
        content = truncate_with_marker(&content, cap, &mut truncated);
    }
    (content, truncated)
}

fn truncate_with_marker(text: &str, max_bytes: usize, truncated: &mut bool) -> String {
    if text.len() <= max_bytes {
        return text.to_string();
    }
    *truncated = true;
    if max_bytes <= TRUNCATION_MARKER.len() {
        return truncate_to_char_boundary(TRUNCATION_MARKER, max_bytes).to_string();
    }
    let text_budget = max_bytes - TRUNCATION_MARKER.len();
    format!(
        "{}{}",
        truncate_to_char_boundary(text, text_budget),
        TRUNCATION_MARKER
    )
}

fn truncate_to_char_boundary(text: &str, max_bytes: usize) -> &str {
    if text.len() <= max_bytes {
        return text;
    }
    let mut end = max_bytes;
    while end > 0 && !text.is_char_boundary(end) {
        end -= 1;
    }
    &text[..end]
}

pub(super) fn validate_public_url(
    url: &Url,
    cancellation: &AgentCancellation,
) -> anyhow::Result<()> {
    cancellation.check()?;
    resolve_public_target(url, cancellation).map(|_| ())
}

fn resolve_public_target(
    url: &Url,
    cancellation: &AgentCancellation,
) -> anyhow::Result<PinnedTarget> {
    let (receiver, worker) = spawn_dns_resolution_worker(url.clone())?;
    wait_for_worker_result(
        receiver,
        worker,
        "url_fetch DNS resolution timeout",
        URL_FETCH_DNS_TIMEOUT,
        cancellation,
    )
}

fn resolve_public_target_sync(url: &Url) -> anyhow::Result<PinnedTarget> {
    if !matches!(url.scheme(), "http" | "https") {
        anyhow::bail!("url_fetch only supports http and https URLs");
    }
    let host = url
        .host_str()
        .ok_or_else(|| anyhow::anyhow!("url_fetch URL must include a host"))?
        .to_string();
    let port = url.port_or_known_default().unwrap_or(80);
    let addrs = resolve_host_addrs(&host, port)?;
    validate_public_addrs(&host, addrs)
}

fn resolve_host_addrs(host: &str, port: u16) -> anyhow::Result<Vec<SocketAddr>> {
    if let Ok(ip) = host.parse::<IpAddr>() {
        return Ok(vec![SocketAddr::new(ip, port)]);
    }
    (host, port)
        .to_socket_addrs()
        .map(|iter| iter.collect::<Vec<_>>())
        .map_err(|error| anyhow::anyhow!("url_fetch DNS resolution failed for {host}: {error}"))
}

fn validate_public_addrs(host: &str, addrs: Vec<SocketAddr>) -> anyhow::Result<PinnedTarget> {
    if addrs.is_empty() {
        anyhow::bail!("url_fetch DNS resolution returned no addresses for {host}");
    }
    for addr in &addrs {
        if !is_public_ip(addr.ip()) {
            anyhow::bail!(
                "url_fetch target resolves to unsafe IP address: {}",
                addr.ip()
            );
        }
    }
    Ok(PinnedTarget {
        host: host.to_string(),
        addrs,
    })
}

fn is_public_ipv4(ip: std::net::Ipv4Addr) -> bool {
    let value = u32::from(ip);
    !(value & 0xff00_0000 == 0x0000_0000 // 0/8
        || value & 0xff00_0000 == 0x0a00_0000 // 10/8 private
        || value & 0xfff0_0000 == 0xac10_0000 // 172.16/12 private
        || value & 0xffff_0000 == 0xc0a8_0000 // 192.168/16 private
        || value & 0xff00_0000 == 0x7f00_0000 // 127/8 loopback
        || value & 0xffff_0000 == 0xa9fe_0000 // 169.254/16 link-local
        || value & 0xffc0_0000 == 0x6440_0000 // 100.64/10 shared
        || value & 0xffff_ff00 == 0xc000_0000 // 192.0.0/24 protocol assignments
        || value & 0xffff_ff00 == 0xc000_0200 // 192.0.2/24 documentation
        || value & 0xffff_ff00 == 0xc633_6400 // 198.51.100/24 documentation
        || value & 0xffff_ff00 == 0xcb00_7100 // 203.0.113/24 documentation
        || value & 0xffff_ff00 == 0xc058_6300 // 192.88.99/24 6to4 anycast
        || value & 0xfffe_0000 == 0xc612_0000 // 198.18/15 benchmarking
        || value & 0xf000_0000 == 0xe000_0000 // multicast
        || value & 0xf000_0000 == 0xf000_0000) // reserved, including broadcast
}

fn is_public_ip(ip: IpAddr) -> bool {
    match ip {
        IpAddr::V4(ip) => is_public_ipv4(ip),
        IpAddr::V6(ip) => {
            if let Some(mapped) = ip.to_ipv4_mapped() {
                return is_public_ipv4(mapped);
            }
            let segments = ip.segments();
            let compatible = segments[..6].iter().all(|segment| *segment == 0);
            let first = segments[0];
            !(
                ip.is_loopback()
                || ip.is_unspecified()
                || ip.is_multicast()
                || compatible // IPv4-compatible ::/96
                || (first & 0xfe00) == 0xfc00 // fc00::/7 ULA
                || (first & 0xffc0) == 0xfe80 // fe80::/10 link-local
                || (first & 0xffc0) == 0xfec0 // fec0::/10 deprecated site-local
                || first == 0x2002 // 2002::/16 6to4
                || (first == 0x2001 && segments[1] <= 0x01ff) // 2001::/23 protocol assignments
                || (first == 0x2001 && segments[1] == 0x0db8) // 2001:db8::/32 documentation
                || (first == 0x3fff && segments[1] <= 0x0fff) // 3fff::/20 documentation
                || first == 0x5f00 // 5f00::/16 special use
                || (first == 0x0064
                    && segments[1] == 0xff9b
                    && (segments[2..6].iter().all(|segment| *segment == 0)
                        || segments[2] == 1)) // 64:ff9b::/96 and 64:ff9b:1::/48
                || (first == 0x0100
                    && segments[1..4].iter().all(|segment| *segment == 0))
                // 100::/64 discard-only
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cancellation::AgentCancellationHandle;
    use std::{
        io::{Read, Write},
        net::{Ipv4Addr, TcpListener},
        sync::{
            atomic::{AtomicUsize, Ordering},
            mpsc,
        },
        thread,
    };

    fn http_response(content_type: &str, body: &str) -> String {
        format!(
            "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\n\r\n{body}",
            body.len()
        )
    }

    fn fake_server(
        response: String,
    ) -> (SocketAddr, mpsc::Receiver<String>, thread::JoinHandle<()>) {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let (tx, rx) = mpsc::channel();
        let handle = thread::spawn(move || {
            let (mut stream, _) = listener.accept().unwrap();
            let mut buf = [0_u8; 8192];
            let n = stream.read(&mut buf).unwrap();
            let _ = tx.send(String::from_utf8_lossy(&buf[..n]).to_string());
            let _ = stream.write_all(response.as_bytes());
        });
        (addr, rx, handle)
    }

    fn pinned_loopback_client(addr: SocketAddr) -> Client {
        build_pinned_client(&PinnedTarget {
            host: "example.test".to_string(),
            addrs: vec![addr],
        })
        .unwrap()
    }

    #[test]
    fn public_ip_policy_rejects_reserved_and_special_ranges() {
        let blocked_ipv4 = [
            [0, 0, 0, 1],
            [10, 0, 0, 1],
            [100, 64, 0, 1],
            [100, 127, 255, 254],
            [127, 0, 0, 1],
            [169, 254, 169, 254],
            [172, 16, 0, 1],
            [192, 0, 0, 1],
            [192, 0, 2, 1],
            [192, 88, 99, 1],
            [192, 168, 0, 1],
            [198, 18, 0, 1],
            [198, 19, 255, 254],
            [198, 51, 100, 1],
            [203, 0, 113, 1],
            [224, 0, 0, 1],
            [240, 0, 0, 1],
            [255, 255, 255, 255],
        ];
        for octets in blocked_ipv4 {
            let ip = IpAddr::V4(Ipv4Addr::from(octets));
            assert!(!is_public_ip(ip), "{ip}");
            assert!(validate_public_addrs("bad", vec![SocketAddr::new(ip, 80)]).is_err());
        }

        let blocked_ipv6 = [
            "::",
            "::1",
            "::8.8.8.8",
            "::ffff:10.0.0.1",
            "64:ff9b::1",
            "64:ff9b:1::1",
            "100::1",
            "2001::1",
            "2001:1::1",
            "2001:db8::1",
            "2002::1",
            "3fff::1",
            "5f00::1",
            "fc00::1",
            "fd00::1",
            "fec0::1",
            "fe80::1",
            "ff02::1",
        ];
        for value in blocked_ipv6 {
            let ip = IpAddr::V6(value.parse().unwrap());
            assert!(!is_public_ip(ip), "{ip}");
            assert!(validate_public_addrs("bad", vec![SocketAddr::new(ip, 80)]).is_err());
        }

        for ip in [
            IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)),
            IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)),
            IpAddr::V4(Ipv4Addr::new(100, 128, 0, 1)),
            IpAddr::V4(Ipv4Addr::new(198, 20, 0, 1)),
            IpAddr::V6("2001:4860:4860::8888".parse().unwrap()),
            IpAddr::V6("2606:4700:4700::1111".parse().unwrap()),
            IpAddr::V6("::ffff:8.8.8.8".parse().unwrap()),
        ] {
            assert!(is_public_ip(ip), "{ip}");
        }

        let cancellation = AgentCancellation::default();
        assert!(
            resolve_public_target(&Url::parse("http://127.0.0.1/").unwrap(), &cancellation)
                .is_err()
        );
        assert!(
            resolve_public_target(&Url::parse("http://localhost/").unwrap(), &cancellation)
                .is_err()
        );
    }

    #[test]
    fn recv_cancellable_is_bounded_and_cancellation_aware() {
        let (_sender, receiver) = mpsc::sync_channel::<()>(1);
        let started = Instant::now();
        let error = recv_cancellable(
            &receiver,
            "test timeout",
            Duration::from_millis(40),
            &AgentCancellation::default(),
        )
        .unwrap_err();
        assert!(error.to_string().contains("test timeout"));
        assert!(started.elapsed() < Duration::from_secs(1));

        let (cancellation, handle) = AgentCancellation::default().child_token();
        handle.cancel();
        let error = recv_cancellable(
            &receiver,
            "test timeout",
            Duration::from_secs(5),
            &cancellation,
        )
        .unwrap_err();
        assert_eq!(error.to_string(), "prompt canceled");
    }

    #[test]
    fn worker_wait_detaches_without_cleanup_grace_after_deadline() {
        let (_result_sender, result_receiver) = mpsc::sync_channel::<anyhow::Result<()>>(1);
        let (done_sender, done_receiver) = mpsc::sync_channel(1);
        let (ready_sender, ready_receiver) = mpsc::channel();
        let (release_sender, release_receiver) = mpsc::channel();
        let (finished_sender, finished_receiver) = mpsc::channel();
        let join_handle = thread::spawn(move || {
            done_sender.send(()).unwrap();
            ready_sender.send(()).unwrap();
            let _ = release_receiver.recv();
            let _ = finished_sender.send(());
        });
        ready_receiver.recv_timeout(Duration::from_secs(1)).unwrap();
        let worker = UrlFetchWorkerHandle {
            worker_label: "test",
            done_receiver,
            join_handle,
        };
        let (returned_sender, returned_receiver) = mpsc::channel();
        let wait_handle = thread::spawn(move || {
            let result = wait_for_worker_result(
                result_receiver,
                worker,
                "test timeout",
                Duration::from_millis(20),
                &AgentCancellation::default(),
            );
            returned_sender.send(result).unwrap();
        });
        let returned = returned_receiver.recv_timeout(Duration::from_secs(1));

        release_sender.send(()).unwrap();
        wait_handle.join().unwrap();
        finished_receiver
            .recv_timeout(Duration::from_secs(1))
            .unwrap();

        let error = returned
            .expect("deadline cleanup should not wait for a worker join")
            .unwrap_err();
        assert!(error.to_string().contains("test timeout"));
    }

    #[test]
    fn dns_resolver_permits_exhaust_and_release_deterministically() {
        let count = AtomicUsize::new(0);
        let first = acquire_dns_resolver_permit(&count, 2).unwrap();
        let second = acquire_dns_resolver_permit(&count, 2).unwrap();
        assert_eq!(count.load(Ordering::Relaxed), 2);

        let error = acquire_dns_resolver_permit(&count, 2)
            .err()
            .expect("the local resolver cap should be exhausted");
        assert_eq!(
            error.to_string(),
            "url_fetch DNS resolver capacity exhausted; retry later"
        );

        drop(first);
        let replacement = acquire_dns_resolver_permit(&count, 2).unwrap();
        assert_eq!(count.load(Ordering::Relaxed), 2);

        drop(second);
        drop(replacement);
        assert_eq!(count.load(Ordering::Relaxed), 0);

        let final_permit = acquire_dns_resolver_permit(&count, 2).unwrap();
        assert_eq!(count.load(Ordering::Relaxed), 1);
        drop(final_permit);
        assert_eq!(count.load(Ordering::Relaxed), 0);
    }
    #[test]
    fn pinned_resolution_is_used_for_request() {
        let (addr, rx, handle) = fake_server(http_response("text/plain", "ok"));
        let client = pinned_loopback_client(addr);
        let args = UrlFetchInput {
            url: "http://example.test/hello".to_string(),
            max_tokens: None,
        };
        let result = fetch_with_client(
            args,
            Url::parse("http://example.test/hello").unwrap(),
            &client,
            &AgentCancellation::default(),
        )
        .unwrap();
        handle.join().unwrap();
        let request = rx.recv().unwrap();
        assert!(request.starts_with("GET /hello HTTP/1.1"), "{request}");
        assert!(
            request.to_ascii_lowercase().contains("host: example.test"),
            "{request}"
        );
        assert!(result.success);
        assert_eq!(result.content, "ok");
    }

    #[test]
    fn response_body_decoding_honors_bom_header_meta_and_utf8_fallback() {
        assert_eq!(
            decode_response_body(b"<p>caf\xe9</p>", "text/html; charset=windows-1252"),
            "<p>café</p>"
        );
        assert_eq!(
            decode_response_body(b"<p>caf\xe9</p>", "text/html; charset=ISO-8859-1"),
            "<p>café</p>"
        );

        let (shift_jis, _, _) = encoding_rs::SHIFT_JIS.encode("<p>日本語</p>");
        assert_eq!(
            decode_response_body(shift_jis.as_ref(), "text/html; charset=Shift_JIS"),
            "<p>日本語</p>"
        );

        let mut bom = vec![0xef, 0xbb, 0xbf];
        bom.extend_from_slice("<p>café</p>".as_bytes());
        assert_eq!(
            decode_response_body(&bom, "text/html; charset=windows-1252"),
            "<p>café</p>"
        );

        assert_eq!(
            decode_response_body(b"<meta charset=windows-1252><p>caf\xe9</p>", "text/html"),
            "<meta charset=windows-1252><p>café</p>"
        );
        assert_eq!(
            decode_response_body(b"<p>caf\xe9</p>", "text/html; charset=unknown-label"),
            "<p>caf�</p>"
        );
    }

    #[test]
    fn redirects_are_not_followed() {
        let response = "HTTP/1.1 302 Found\r\nLocation: http://example.test/final\r\nContent-Length: 8\r\n\r\nredirect".to_string();
        let (addr, rx, handle) = fake_server(response);
        let client = pinned_loopback_client(addr);
        let result = fetch_with_client(
            UrlFetchInput {
                url: "http://example.test/redirect".to_string(),
                max_tokens: None,
            },
            Url::parse("http://example.test/redirect").unwrap(),
            &client,
            &AgentCancellation::default(),
        )
        .unwrap();
        handle.join().unwrap();
        let request = rx.recv().unwrap();
        assert!(request.starts_with("GET /redirect HTTP/1.1"), "{request}");
        assert!(!result.success);
        assert_eq!(result.metadata[meta::STATUS], 302);
    }

    #[test]
    fn content_length_over_max_bytes_fails_before_reading_body() {
        let response = format!(
            "HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\nsmall",
            URL_FETCH_MAX_BYTES + 1
        );
        let (addr, _, handle) = fake_server(response);
        let client = pinned_loopback_client(addr);
        let error = fetch_response_with_client(
            &client,
            &Url::parse("http://example.test/too-large").unwrap(),
            &AgentCancellation::default(),
        )
        .unwrap_err();
        handle.join().unwrap();
        assert!(error.to_string().contains("response exceeded"), "{error}");
    }

    #[test]
    fn streamed_body_over_max_bytes_fails_without_content_length() {
        let response = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nConnection: close\r\n\r\n{}",
            "x".repeat(URL_FETCH_MAX_BYTES + 1)
        );
        let (addr, _, handle) = fake_server(response);
        let client = pinned_loopback_client(addr);
        let error = fetch_response_with_client(
            &client,
            &Url::parse("http://example.test/stream-too-large").unwrap(),
            &AgentCancellation::default(),
        )
        .unwrap_err();
        handle.join().unwrap();
        assert!(error.to_string().contains("response exceeded"), "{error}");
    }

    #[test]
    fn output_truncation_preserves_utf8_char_boundaries() {
        let (content, truncated) = sanitize_and_bound(&"界".repeat(2_000), 1_024);
        assert!(truncated);
        assert!(content.ends_with(TRUNCATION_MARKER));
        assert!(content.len() <= 1_024 * 4);
        assert!(content.is_char_boundary(content.len() - TRUNCATION_MARKER.len()));
    }

    #[test]
    fn non_success_status_returns_bounded_redacted_body() {
        let secret = "sk-webFetchSecret123456";
        let body = format!(
            "failure token={secret} {}",
            "x".repeat(URL_FETCH_OUTPUT_MAX_BYTES)
        );
        let response = format!(
            "HTTP/1.1 500 Internal Server Error\r\nContent-Type: text/plain\r\nContent-Length: {}\r\n\r\n{}",
            body.len(),
            body
        );
        let (addr, _, handle) = fake_server(response);
        let client = pinned_loopback_client(addr);
        let result = fetch_with_client(
            UrlFetchInput {
                url: "http://example.test/error".to_string(),
                max_tokens: Some(10_000),
            },
            Url::parse("http://example.test/error").unwrap(),
            &client,
            &AgentCancellation::default(),
        )
        .unwrap();
        handle.join().unwrap();
        assert!(!result.success);
        assert_eq!(result.metadata[meta::STATUS], 500);
        assert_eq!(result.metadata[meta::TRUNCATED], true);
        assert!(result.content.len() <= URL_FETCH_OUTPUT_MAX_BYTES);
        assert!(!result.content.contains(secret));
        assert!(result.content.contains("<redacted>"));
    }

    #[test]
    fn html_uses_readability_when_article_is_long_enough() {
        let html = format!(
            "<html><head><title>Noise</title></head><body><article><h1>Readable Title</h1><p>{}</p></article></body></html>",
            "Readable article sentence. ".repeat(20)
        );
        let extracted = extract_content(&html, "text/html", "https://example.com/page");
        assert!(
            extracted.text.contains("Readable article sentence"),
            "{}",
            extracted.text
        );
        assert!(extracted.title.is_some());
    }

    #[test]
    fn short_html_falls_back_to_fast_html2md() {
        let extracted = extract_content(
            "<html><body><p>Short <strong>fallback</strong></p></body></html>",
            "text/html; charset=utf-8",
            "https://example.com",
        );
        assert!(extracted.text.contains("Short"));
        assert!(extracted.text.contains("fallback"));
        assert_eq!(extracted.title, None);
    }

    #[test]
    fn non_html_returns_raw_redacted_text_and_truncates() {
        let secret = "sk-webFetchSecret123456";
        let (content, truncated) = sanitize_and_bound(
            &format!(
                "{{\"token\":\"{secret}\"}}{}",
                "x".repeat(URL_FETCH_OUTPUT_MAX_BYTES)
            ),
            10_000,
        );
        assert!(truncated);
        assert!(content.ends_with(TRUNCATION_MARKER));
        assert!(!content.contains(secret));
        assert!(content.contains("<redacted>"));
        assert!(content.len() <= URL_FETCH_OUTPUT_MAX_BYTES);

        let extracted = extract_content("{\"ok\":true}", "application/json", "https://example.com");
        assert_eq!(extracted.text, "{\"ok\":true}");
    }

    #[test]
    fn canceled_during_blocking_body_read_returns_without_waiting_for_request_timeout() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let (release_tx, release_rx) = mpsc::channel::<()>();
        let server = thread::spawn(move || {
            let (mut stream, _) = listener.accept().unwrap();
            let mut buf = [0_u8; 8192];
            let _ = stream.read(&mut buf).unwrap();
            let _ = stream.write_all(
                b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: 4\r\n\r\n",
            );
            let _ = release_rx.recv();
        });
        let client = pinned_loopback_client(addr);
        let (cancellation, handle): (AgentCancellation, AgentCancellationHandle) =
            AgentCancellation::default().child_token();
        let canceler = thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            handle.cancel();
        });
        let start = Instant::now();

        let error = fetch_response_with_client(
            &client,
            &Url::parse("http://example.test/hang").unwrap(),
            &cancellation,
        )
        .unwrap_err();

        release_tx.send(()).unwrap();
        canceler.join().unwrap();
        server.join().unwrap();
        assert_eq!(error.to_string(), "prompt canceled");
        assert!(start.elapsed() < Duration::from_secs(2));
    }
}