//! `citation` matcher: fetch a response's cited sources and judge grounding.
//!
//! Why this exists: a model can state a claim and cite a URL that does not
//! exist, or that says the opposite. `contains`/`regex` see the citation is
//! present; they cannot see whether it backs the claim. The `citation` matcher
//! extracts the URLs a response cites, fetches each one behind an SSRF guard,
//! and asks an LLM judge whether the fetched source actually supports the
//! claim. It passes when enough cited sources support the claim.
//!
//! Like `similar` and `llm-judge`, this is an async matcher: it needs a
//! provider call (and a network fetch). It runs through the executor's async
//! path, not the synchronous `check_matcher` dispatch, and it is not composable
//! under `not:`.
//!
//! Network fetch is the only impure step, so it goes through a [`SourceFetcher`]
//! trait. Production wires [`crate::matchers::citation_fetch::ReqwestFetcher`];
//! tests inject a map so the matcher, the SSRF classification, and the grounding
//! fold are all exercised offline.

use std::net::IpAddr;
use std::str::FromStr;

use async_trait::async_trait;
use serde::Deserialize;
use serde_json::Value;

use super::llm_judge::{self, JudgeConfig};
use super::result::{AssertionDiff, AssertionResult, MatcherKind};
use crate::eval::provider::LlmProvider;
use crate::eval::ProviderError;
use crate::Error;

/// Parsed `citation:` YAML block.
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct CitationConfig {
    /// The claim each cited source is judged against.
    pub claim: String,
    /// Judge model identifier passed to the provider.
    pub model: String,
    /// Fraction of cited sources in `[0, 1]` that must support the claim for
    /// the assertion to pass. Defaults to 1.0: every citation must hold up. A
    /// blocked or unreachable citation counts as unsupported, so a fabricated
    /// URL drags the ratio down.
    #[serde(default = "default_threshold")]
    pub threshold: f64,
    /// Fail when the response cites no fetchable source. Defaults to true: a
    /// claim that grounds itself in nothing is not grounded.
    #[serde(default = "default_true")]
    pub require_citation: bool,
}

fn default_threshold() -> f64 {
    1.0
}

fn default_true() -> bool {
    true
}

/// Parse a raw matcher value into a typed [`CitationConfig`].
///
/// Returns [`Error::InvalidMatcherInput`] when the block does not match the
/// schema or the threshold is out of range, so the executor folds it like every
/// other matcher error.
pub fn parse_config(value: &Value) -> crate::Result<CitationConfig> {
    let config: CitationConfig = serde_json::from_value(value.clone())
        .map_err(|e| Error::InvalidMatcherInput(format!("invalid citation block: {e}")))?;
    if !(0.0..=1.0).contains(&config.threshold) {
        return Err(Error::InvalidMatcherInput(
            "citation `threshold` must be in [0, 1]".into(),
        ));
    }
    Ok(config)
}

/// Why a citation could not be fetched.
#[derive(Debug)]
pub enum FetchError {
    /// The SSRF guard rejected the URL (private/loopback/metadata host, or a
    /// non-http scheme) before any connection was made.
    Blocked(String),
    /// The source could not be retrieved (DNS, connection, status, or size).
    Unreachable(String),
}

impl std::fmt::Display for FetchError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            FetchError::Blocked(why) => write!(f, "blocked: {why}"),
            FetchError::Unreachable(why) => write!(f, "unreachable: {why}"),
        }
    }
}

/// Fetches the text of a cited source. Injectable so the matcher is testable
/// offline without a network call.
#[async_trait]
pub trait SourceFetcher: Send + Sync {
    /// Fetch the source at `url`, returning its text body.
    async fn fetch(&self, url: &str) -> Result<String, FetchError>;
}

/// Extract the http(s) URLs a value cites, in first-seen order with duplicates
/// removed. Walks every string in the JSON value (text, nested objects, arrays)
/// so a citation in a tool result or an agent's final response is found the
/// same way.
pub fn extract_urls(value: &Value) -> Vec<String> {
    let mut out = Vec::new();
    collect_urls(value, &mut out);
    out
}

fn collect_urls(value: &Value, out: &mut Vec<String>) {
    match value {
        Value::String(s) => urls_in_text(s, out),
        Value::Array(items) => items.iter().for_each(|v| collect_urls(v, out)),
        Value::Object(map) => map.values().for_each(|v| collect_urls(v, out)),
        _ => {}
    }
}

/// Pull every `http://` or `https://` run out of one string. A URL ends at the
/// first whitespace or character that cannot appear in one; trailing sentence
/// punctuation is trimmed so "see https://example.com." cites the bare URL.
fn urls_in_text(text: &str, out: &mut Vec<String>) {
    let mut rest = text;
    while let Some(idx) = find_scheme(rest) {
        let tail = &rest[idx..];
        let end = tail
            .find(|c: char| {
                c.is_whitespace() || matches!(c, '"' | '<' | '>' | ')' | ']' | '}' | '|')
            })
            .unwrap_or(tail.len());
        let url = tail[..end].trim_end_matches(['.', ',', ';', ':', '!', '?', '\'']);
        if !url.is_empty() && !out.iter().any(|u| u == url) {
            out.push(url.to_string());
        }
        rest = &tail[end..];
    }
}

/// Byte index of the next `http://` or `https://` in `text`, or `None`.
fn find_scheme(text: &str) -> Option<usize> {
    let http = text.find("http://");
    let https = text.find("https://");
    match (http, https) {
        (Some(a), Some(b)) => Some(a.min(b)),
        (Some(a), None) => Some(a),
        (None, b) => b,
    }
}

/// Classify a URL for SSRF safety before any fetch. Rejects a non-http scheme
/// and a host that is a private, loopback, link-local (cloud-metadata), or
/// unspecified IP literal, plus the `localhost` name. A hostname that resolves
/// to a blocked IP is caught later by the fetcher after DNS; this is the
/// pre-connection gate.
pub fn classify_url(url: &str) -> Result<(), String> {
    let lowered = url.to_ascii_lowercase();
    if !lowered.starts_with("http://") && !lowered.starts_with("https://") {
        return Err("only http(s) URLs may be fetched".to_string());
    }
    let Some(host) = url_host(url) else {
        return Err("URL has no host".to_string());
    };
    if host.eq_ignore_ascii_case("localhost") {
        return Err("host `localhost` is not fetchable".to_string());
    }
    if let Ok(ip) = IpAddr::from_str(host) {
        if is_blocked_ip(&ip) {
            return Err(format!("host {host} is a private or metadata address"));
        }
    }
    Ok(())
}

/// Extract the host from a URL without a full parser. Handles
/// `scheme://user@host:port/path` and bracketed IPv6 literals.
fn url_host(url: &str) -> Option<&str> {
    let authority = url.split("://").nth(1)?.split('/').next()?;
    let host_port = authority.rsplit('@').next()?;
    if let Some(rest) = host_port.strip_prefix('[') {
        return rest.split(']').next();
    }
    Some(host_port.split(':').next().unwrap_or(host_port))
}

/// True for an IP a fetch must never reach: private, loopback, link-local
/// (which covers the 169.254.169.254 metadata address), or unspecified.
pub fn is_blocked_ip(ip: &IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => {
            v4.is_private() || v4.is_loopback() || v4.is_link_local() || v4.is_unspecified()
        }
        IpAddr::V6(v6) => v6.is_loopback() || v6.is_unspecified() || is_v6_unique_local(v6),
    }
}

/// IPv6 unique-local (fc00::/7), the v6 analogue of RFC 1918. `Ipv6Addr` has no
/// stable `is_unique_local` on this toolchain, so check the prefix directly.
fn is_v6_unique_local(v6: &std::net::Ipv6Addr) -> bool {
    (v6.segments()[0] & 0xfe00) == 0xfc00
}

/// One cited source's grounding outcome.
#[derive(Debug, Clone, PartialEq, Eq)]
enum SourceOutcome {
    Supported,
    Unsupported,
    Blocked,
    Unreachable,
}

/// Judge a response's cited sources against the claim and fold the verdicts.
///
/// Extracts the cited URLs, drops the SSRF-blocked ones, fetches the rest
/// through `fetcher`, and asks `provider` (one judge call per source) whether
/// the fetched text supports `config.claim`. Passes when the supported fraction
/// clears `config.threshold`. A blocked or unreachable citation counts against
/// the fraction, so a fabricated URL fails the assertion. A provider error
/// propagates as a [`ProviderError`].
pub async fn evaluate(
    config: &CitationConfig,
    actual: &Value,
    provider: &dyn LlmProvider,
    fetcher: &dyn SourceFetcher,
) -> Result<AssertionResult, ProviderError> {
    let urls = extract_urls(actual);
    if urls.is_empty() {
        return Ok(no_citations_result(config));
    }
    let mut outcomes: Vec<(String, SourceOutcome)> = Vec::with_capacity(urls.len());
    for url in &urls {
        if let Err(reason) = classify_url(url) {
            outcomes.push((format!("{url}: blocked ({reason})"), SourceOutcome::Blocked));
            continue;
        }
        let source = match fetcher.fetch(url).await {
            Ok(text) => text,
            Err(err) => {
                let kind = match err {
                    FetchError::Blocked(_) => SourceOutcome::Blocked,
                    FetchError::Unreachable(_) => SourceOutcome::Unreachable,
                };
                outcomes.push((format!("{url}: {err}"), kind));
                continue;
            }
        };
        let judged = judge_source(config, &source, provider).await?;
        let outcome = if judged {
            SourceOutcome::Supported
        } else {
            SourceOutcome::Unsupported
        };
        let label = if judged {
            "supports"
        } else {
            "does not support"
        };
        outcomes.push((format!("{url}: {label} the claim"), outcome));
    }
    Ok(fold(config, &outcomes))
}

/// Ask the judge whether one source supports the claim.
async fn judge_source(
    config: &CitationConfig,
    source: &str,
    provider: &dyn LlmProvider,
) -> Result<bool, ProviderError> {
    let rubric = format!(
        "You are checking whether a SOURCE supports a CLAIM.\n\nCLAIM: {}\n\nPass only if \
         the source's content explicitly supports the claim. Fail if it contradicts the claim, \
         is unrelated, or does not mention it.",
        config.claim
    );
    let judge_cfg = JudgeConfig {
        rubric,
        threshold: 0.5,
        model: Some(config.model.clone()),
    };
    let result = llm_judge::evaluate(
        &judge_cfg,
        &Value::String(source.to_string()),
        provider,
        &config.model,
    )
    .await?;
    Ok(result.pass)
}

/// Build the result for a response that cites nothing.
fn no_citations_result(config: &CitationConfig) -> AssertionResult {
    if config.require_citation {
        AssertionResult::failed(
            MatcherKind::Citation,
            AssertionDiff {
                path: String::new(),
                expected: None,
                actual: None,
                note: Some("no fetchable citation found in the response".to_string()),
            },
        )
    } else {
        AssertionResult::passed_with_note(
            MatcherKind::Citation,
            "no citation to verify (require_citation is false)".to_string(),
        )
    }
}

/// Fold per-source outcomes into the assertion result against the threshold.
fn fold(config: &CitationConfig, outcomes: &[(String, SourceOutcome)]) -> AssertionResult {
    let total = outcomes.len();
    let supported = outcomes
        .iter()
        .filter(|(_, o)| *o == SourceOutcome::Supported)
        .count();
    let ratio = if total > 0 {
        supported as f64 / total as f64
    } else {
        0.0
    };
    let pass = ratio >= config.threshold;
    let summary = format!(
        "citation: {supported}/{total} cited source(s) support the claim (ratio {ratio:.2} \
         vs threshold {:.2})",
        config.threshold
    );
    let mut diffs = vec![AssertionDiff {
        path: String::new(),
        expected: Some(Value::from(config.threshold)),
        actual: Some(Value::from(ratio)),
        note: Some(summary),
    }];
    for (note, _) in outcomes {
        diffs.push(AssertionDiff {
            path: String::new(),
            expected: None,
            actual: None,
            note: Some(note.clone()),
        });
    }
    AssertionResult {
        kind: MatcherKind::Citation,
        pass,
        diffs,
        captures: Default::default(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    use async_trait::async_trait;
    use serde_json::json;

    use crate::eval::provider::{
        Capabilities, CompletionOptions, Cost, FinishReason, Message, ProviderId, Response,
        ResponseContent, Usage,
    };

    // -- pure helpers --------------------------------------------------------

    #[test]
    fn extract_urls_finds_and_dedups_across_a_nested_value() {
        let v = json!({
            "text": "See https://example.com/a and https://example.com/a again.",
            "refs": ["http://docs.example.org/b)", "not a url"],
        });
        // Dedups the repeat and ignores the non-URL string. Object-key
        // iteration order is not part of the contract, so compare as a set.
        let mut urls = extract_urls(&v);
        urls.sort();
        assert_eq!(
            urls,
            vec![
                "http://docs.example.org/b".to_string(),
                "https://example.com/a".to_string(),
            ]
        );
    }

    #[test]
    fn extract_urls_trims_trailing_sentence_punctuation() {
        let urls = extract_urls(&json!("source: https://example.com/page."));
        assert_eq!(urls, vec!["https://example.com/page".to_string()]);
    }

    #[test]
    fn classify_url_blocks_metadata_loopback_and_non_http() {
        assert!(classify_url("http://169.254.169.254/").is_err());
        assert!(classify_url("http://localhost/x").is_err());
        assert!(classify_url("http://127.0.0.1/x").is_err());
        assert!(classify_url("http://10.0.0.1/x").is_err());
        assert!(classify_url("ftp://example.com/x").is_err());
        assert!(classify_url("https://example.com/ok").is_ok());
    }

    #[test]
    fn parse_config_rejects_out_of_range_threshold() {
        let v = json!({"claim": "x", "model": "m", "threshold": 2.0});
        assert!(matches!(
            parse_config(&v),
            Err(Error::InvalidMatcherInput(_))
        ));
    }

    #[test]
    fn parse_config_defaults_threshold_and_require_citation() {
        let cfg = parse_config(&json!({"claim": "x", "model": "m"})).unwrap();
        assert_eq!(cfg.threshold, 1.0);
        assert!(cfg.require_citation);
    }

    // -- the matcher, offline ------------------------------------------------

    /// Returns the mapped body for a known URL, or unreachable otherwise.
    struct MapFetcher(HashMap<String, String>);

    #[async_trait]
    impl SourceFetcher for MapFetcher {
        async fn fetch(&self, url: &str) -> Result<String, FetchError> {
            self.0
                .get(url)
                .cloned()
                .ok_or_else(|| FetchError::Unreachable("404".to_string()))
        }
    }

    /// A judge that passes only when the source text contains a marker word,
    /// standing in for "the source supports the claim".
    struct MarkerJudge {
        marker: &'static str,
    }

    #[async_trait]
    impl LlmProvider for MarkerJudge {
        fn name(&self) -> ProviderId {
            ProviderId::Anthropic
        }
        fn supported_capabilities(&self) -> Capabilities {
            Capabilities::none()
        }
        async fn chat_completion(
            &self,
            messages: &[Message],
            _options: &CompletionOptions,
        ) -> Result<Response, ProviderError> {
            // The judge sees the source as the candidate in the final message.
            let supports = messages
                .last()
                .map(|m| m.content.contains(self.marker))
                .unwrap_or(false);
            let verdict = if supports {
                r#"{"pass": true, "score": 0.95, "reason": "supports"}"#
            } else {
                r#"{"pass": false, "score": 0.1, "reason": "no support"}"#
            };
            Ok(Response {
                content: ResponseContent::Text {
                    text: verdict.to_string(),
                },
                finish_reason: FinishReason::Stop,
                usage: Usage::default(),
                model: "stub".into(),
            })
        }
        async fn count_tokens(&self, _t: &str, _m: &str) -> Result<u32, ProviderError> {
            Ok(0)
        }
        async fn embed(&self, _t: &str, _m: &str) -> Result<Vec<f32>, ProviderError> {
            Ok(vec![1.0])
        }
        fn estimate_cost(&self, _u: &Usage, _m: &str) -> Cost {
            Cost::default()
        }
    }

    fn config(threshold: f64) -> CitationConfig {
        CitationConfig {
            claim: "the sky is blue".into(),
            model: "judge".into(),
            threshold,
            require_citation: true,
        }
    }

    #[tokio::test]
    async fn passes_when_the_cited_source_supports_the_claim() {
        let fetcher = MapFetcher(HashMap::from([(
            "https://example.com/sky".to_string(),
            "Studies confirm this: GROUNDED, the sky is blue from Rayleigh scattering.".to_string(),
        )]));
        let actual = json!("The sky is blue, see https://example.com/sky");
        let r = evaluate(
            &config(1.0),
            &actual,
            &MarkerJudge { marker: "GROUNDED" },
            &fetcher,
        )
        .await
        .unwrap();
        assert!(r.pass, "diffs: {:?}", r.diffs);
        assert_eq!(r.kind, MatcherKind::Citation);
    }

    #[tokio::test]
    async fn fails_when_the_cited_source_does_not_support_the_claim() {
        let fetcher = MapFetcher(HashMap::from([(
            "https://example.com/grass".to_string(),
            "This page is about how grass is green.".to_string(),
        )]));
        let actual = json!("The sky is blue, see https://example.com/grass");
        let r = evaluate(
            &config(1.0),
            &actual,
            &MarkerJudge { marker: "GROUNDED" },
            &fetcher,
        )
        .await
        .unwrap();
        assert!(!r.pass);
        assert!(r.diffs[0].note.as_deref().unwrap().contains("0/1"));
    }

    #[tokio::test]
    async fn fails_on_a_fabricated_unreachable_citation() {
        // The cited URL is well-formed and passes the SSRF guard, but the
        // fetcher cannot reach it: a fabricated citation fails grounding.
        let fetcher = MapFetcher(HashMap::new());
        let actual = json!("The sky is blue, see https://example.com/does-not-exist");
        let r = evaluate(
            &config(1.0),
            &actual,
            &MarkerJudge { marker: "GROUNDED" },
            &fetcher,
        )
        .await
        .unwrap();
        assert!(!r.pass);
        assert!(r.diffs.iter().any(|d| d
            .note
            .as_deref()
            .map(|n| n.contains("unreachable"))
            .unwrap_or(false)));
    }

    #[tokio::test]
    async fn a_blocked_ssrf_url_counts_against_grounding() {
        let fetcher = MapFetcher(HashMap::new());
        let actual = json!("Internal source: http://169.254.169.254/secret");
        let r = evaluate(
            &config(1.0),
            &actual,
            &MarkerJudge { marker: "GROUNDED" },
            &fetcher,
        )
        .await
        .unwrap();
        assert!(!r.pass);
        assert!(r.diffs.iter().any(|d| d
            .note
            .as_deref()
            .map(|n| n.contains("blocked"))
            .unwrap_or(false)));
    }

    #[tokio::test]
    async fn no_citation_fails_when_required_and_passes_when_not() {
        let fetcher = MapFetcher(HashMap::new());
        let actual = json!("The sky is blue. (no source given)");
        let required = evaluate(
            &config(1.0),
            &actual,
            &MarkerJudge { marker: "GROUNDED" },
            &fetcher,
        )
        .await
        .unwrap();
        assert!(!required.pass);

        let mut cfg = config(1.0);
        cfg.require_citation = false;
        let optional = evaluate(&cfg, &actual, &MarkerJudge { marker: "GROUNDED" }, &fetcher)
            .await
            .unwrap();
        assert!(optional.pass);
    }

    #[tokio::test]
    async fn threshold_allows_partial_support() {
        // Two citations, one supports and one does not: a 0.5 threshold passes,
        // the default 1.0 would fail.
        let fetcher = MapFetcher(HashMap::from([
            (
                "https://example.com/a".to_string(),
                "GROUNDED: the sky is blue".to_string(),
            ),
            (
                "https://example.com/b".to_string(),
                "unrelated content".to_string(),
            ),
        ]));
        let actual = json!("Blue sky: https://example.com/a and https://example.com/b");
        let r = evaluate(
            &config(0.5),
            &actual,
            &MarkerJudge { marker: "GROUNDED" },
            &fetcher,
        )
        .await
        .unwrap();
        assert!(r.pass, "diffs: {:?}", r.diffs);
    }
}
