use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use rand::RngExt as _;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use subtle::ConstantTimeEq as _;
use tracing::debug;

/// Errors that can occur during the OAuth 2.1 + PKCE flow.
#[derive(Debug, thiserror::Error)]
pub enum OAuthError {
    #[error("AS discovery failed for server: {0}")]
    DiscoveryFailed(String),
    #[error("Dynamic Client Registration failed: {0}")]
    DcrFailed(String),
    #[error("Token exchange failed: {0}")]
    TokenExchangeFailed(String),
    #[error("No browser auth available (tunnel not configured or unreachable)")]
    NoBrowserAuth,
    #[error("Invalid OAuth state token — possible CSRF attempt")]
    InvalidState,
    #[error("Missing client_id: server has no registration_endpoint and mcp.json has no clientId")]
    MissingClientId,
    #[error("Missing endpoint: {0}")]
    MissingEndpoint(String),
}

/// In-flight OAuth session stored in bot memory until callback arrives.
pub struct PendingAuth {
    pub server_name: String,
    pub server_url: String,
    pub code_verifier: String,
    pub state: String,
    pub token_endpoint: String,
    pub client_id: String,
    pub client_secret: Option<String>,
    pub redirect_uri: String,
    pub created_at: std::time::Instant,
}

/// Authorization Server Metadata (RFC 8414 / OIDC well-known).
#[derive(Debug, Deserialize)]
pub struct AsMetadata {
    pub authorization_endpoint: String,
    pub token_endpoint: String,
    pub registration_endpoint: Option<String>,
    pub code_challenge_methods_supported: Option<Vec<String>>,
}

/// Resource Metadata (RFC 9728) — points to the Authorization Server.
#[derive(Debug, Deserialize)]
pub struct ResourceMetadata {
    pub authorization_servers: Vec<String>,
    pub scopes_supported: Option<Vec<String>>,
}

/// Dynamic Client Registration response (RFC 7591).
#[derive(Debug, Deserialize)]
pub struct DcrResponse {
    pub client_id: String,
    pub client_secret: Option<String>,
}

/// Token response from the authorization server.
#[derive(Debug, Deserialize)]
pub struct TokenResponse {
    pub access_token: String,
    pub refresh_token: Option<String>,
    pub token_type: Option<String>,
    pub scope: Option<String>,
    pub expires_in: Option<u64>,
}

/// Generate a PKCE code_verifier and code_challenge (S256 method).
///
/// Returns `(code_verifier, code_challenge)` where both are base64url-no-pad encoded.
/// - verifier: 32 random bytes → 43-char base64url string
/// - challenge: SHA-256(verifier_bytes) → 43-char base64url string (per RFC 7636 §4.2)
pub fn generate_pkce() -> (String, String) {
    let mut bytes = [0u8; 32];
    rand::rng().fill(&mut bytes);
    let code_verifier = URL_SAFE_NO_PAD.encode(bytes);
    let hash = Sha256::digest(code_verifier.as_bytes());
    let code_challenge = URL_SAFE_NO_PAD.encode(hash);
    (code_verifier, code_challenge)
}

/// Generate a cryptographically random OAuth state token.
///
/// Returns a 22-char base64url-no-pad string (16 random bytes).
pub fn generate_state() -> String {
    let mut bytes = [0u8; 16];
    rand::rng().fill(&mut bytes);
    URL_SAFE_NO_PAD.encode(bytes)
}

/// Verify an OAuth state token in constant time to prevent timing attacks.
///
/// Per D-05: uses `subtle::ConstantTimeEq` to defeat side-channel attacks.
pub fn verify_state(expected: &str, received: &str) -> bool {
    expected.as_bytes().ct_eq(received.as_bytes()).into()
}

// ---------------------------------------------------------------------------
// Task 1: AS discovery — RFC 9728 → RFC 8414 → OIDC fallback chain
// ---------------------------------------------------------------------------

/// Ordered list of discovery URLs to try for the given server URL.
///
/// Returns URLs in discovery priority order:
/// 1. RFC 9728 resource metadata
/// 2. RFC 8414 AS metadata (no-path)
/// 3. OIDC well-known (no-path)
///
/// This helper is exposed for unit-testing URL construction without HTTP.
pub fn discovery_urls(server_url: &str) -> Vec<String> {
    let parsed = match reqwest::Url::parse(server_url) {
        Ok(u) => u,
        Err(_) => return vec![],
    };
    let origin = format!("{}://{}", parsed.scheme(), parsed.host_str().unwrap_or(""));
    let port_str = parsed.port().map(|p| format!(":{p}")).unwrap_or_default();
    let origin = format!("{origin}{port_str}");

    // Path component (strip trailing slash)
    let raw_path = parsed.path();
    let path = raw_path.trim_end_matches('/');

    let mut urls = Vec::new();

    // Step 1 — RFC 9728 resource metadata
    if path.is_empty() || path == "/" {
        urls.push(format!("{origin}/.well-known/oauth-protected-resource"));
    } else {
        urls.push(format!(
            "{origin}/.well-known/oauth-protected-resource{path}"
        ));
    }

    // Steps 2 & 3 — RFC 8414 and OIDC (no-path variants used here; path variants
    // require the AS URL from the RFC 9728 response, resolved at runtime)
    urls.push(format!("{origin}/.well-known/oauth-authorization-server"));
    urls.push(format!("{origin}/.well-known/openid-configuration"));

    urls
}

/// Build the ordered list of AS metadata URLs for a given AS URL.
///
/// Used in Steps 2 & 3 of `discover_as` after obtaining the AS URL from RFC 9728.
fn as_metadata_urls(as_url: &str) -> Vec<String> {
    let parsed = match reqwest::Url::parse(as_url) {
        Ok(u) => u,
        Err(_) => return vec![],
    };
    let origin = format!("{}://{}", parsed.scheme(), parsed.host_str().unwrap_or(""));
    let port_str = parsed.port().map(|p| format!(":{p}")).unwrap_or_default();
    let origin = format!("{origin}{port_str}");

    let raw_path = parsed.path();
    let path = raw_path.trim_end_matches('/');

    let mut urls = Vec::new();

    if path.is_empty() || path == "/" {
        // No-path AS URL: try RFC 8414, then OIDC.
        urls.push(format!("{origin}/.well-known/oauth-authorization-server"));
        urls.push(format!("{origin}/.well-known/openid-configuration"));
    } else {
        // Path-bearing AS URL: try RFC 8414 §3.1 path-aware variants first…
        urls.push(format!(
            "{origin}/.well-known/oauth-authorization-server{path}"
        ));
        urls.push(format!("{origin}/.well-known/openid-configuration{path}"));
        urls.push(format!("{origin}{path}/.well-known/openid-configuration"));
        // …then origin-only fallbacks. Many real MCP servers (Linear, …) host
        // metadata at the origin-only well-known location even when the
        // server URL itself has a path component.
        urls.push(format!("{origin}/.well-known/oauth-authorization-server"));
        urls.push(format!("{origin}/.well-known/openid-configuration"));
    }

    urls
}

/// Extract `resource_metadata="<url>"` (RFC 9728 §5.1) from a `WWW-Authenticate`
/// header value. Returns `None` if the parameter is absent or malformed.
///
/// Supports both the quoted RFC-compliant form (`resource_metadata="https://…"`)
/// and the unquoted form some servers emit. We deliberately do not parse the
/// full Bearer challenge grammar (RFC 6750 §3) — only the one parameter we need.
/// The parameter-name search is case-sensitive — RFC 9110 §11.2 allows
/// case-insensitive names but real-world servers (Linear, Notion, …)
/// always emit lowercase `resource_metadata`.
fn parse_www_authenticate_resource_metadata(header: &str) -> Option<String> {
    let needle = "resource_metadata=";
    let start = header.find(needle)? + needle.len();
    let rest = &header[start..];
    if let Some(stripped) = rest.strip_prefix('"') {
        let end = stripped.find('"')?;
        let value = &stripped[..end];
        if value.is_empty() {
            None
        } else {
            Some(value.to_string())
        }
    } else {
        let end = rest.find(',').unwrap_or(rest.len());
        let value = rest[..end].trim();
        if value.is_empty() {
            None
        } else {
            Some(value.to_string())
        }
    }
}

/// Step 0 of MCP Authorization-spec discovery: probe the MCP server URL
/// unauthenticated and read RFC 9728 metadata location from the
/// `WWW-Authenticate: Bearer resource_metadata="<url>"` challenge.
///
/// Returns `Some(url)` if the server returned a 401 with a parseable
/// `resource_metadata` parameter. Any other outcome (success, non-401
/// error, missing/malformed header, network error) yields `None` so
/// the caller can fall through to the well-known chain.
/// Only `401` is treated as the WWW-Authenticate signal — `403` and other
/// status codes (including `200`) fall through to the well-known chain.
async fn probe_resource_metadata_via_www_authenticate(
    client: &reqwest::Client,
    server_url: &str,
) -> Option<String> {
    let resp = match client.get(server_url).send().await {
        Ok(r) => r,
        Err(e) => {
            debug!("WWW-Authenticate probe: request failed for {server_url}: {e}");
            return None;
        }
    };
    if resp.status().as_u16() != 401 {
        debug!(
            "WWW-Authenticate probe: {server_url} returned {} (not 401), skipping",
            resp.status()
        );
        return None;
    }
    let header = resp
        .headers()
        .get(reqwest::header::WWW_AUTHENTICATE)?
        .to_str()
        .ok()?;
    let url = parse_www_authenticate_resource_metadata(header)?;
    debug!("WWW-Authenticate probe: resource_metadata = {url}");
    Some(url)
}

/// Discover the Authorization Server metadata for a given MCP server URL.
///
/// Implements the 3-step fallback chain mandated by the MCP Authorization spec:
/// 1. RFC 9728 resource metadata (`/.well-known/oauth-protected-resource`)
/// 2. RFC 8414 AS metadata (`/.well-known/oauth-authorization-server`)
/// 3. OIDC discovery (`/.well-known/openid-configuration`)
///
/// Any non-2xx response (or connection error) on a speculative probe URL
/// is logged and skipped — the URL simply doesn't host metadata. Only when
/// every URL in the chain fails does this return `DiscoveryFailed`.
pub async fn discover_as(
    client: &reqwest::Client,
    server_url: &str,
) -> Result<AsMetadata, OAuthError> {
    let parsed = reqwest::Url::parse(server_url).map_err(|e| {
        OAuthError::DiscoveryFailed(format!("invalid server URL {server_url}: {e}"))
    })?;
    let origin = {
        let scheme = parsed.scheme();
        let host = parsed.host_str().unwrap_or("");
        let port = parsed.port().map(|p| format!(":{p}")).unwrap_or_default();
        format!("{scheme}://{host}{port}")
    };

    let raw_path = parsed.path();
    let path = raw_path.trim_end_matches('/');

    // --- Step 0: WWW-Authenticate probe (MCP Authorization spec / RFC 9728 §5.1) ---
    // Send unauthenticated request to the MCP URL; the server should respond
    // 401 with a `WWW-Authenticate: Bearer resource_metadata="<url>"` header
    // pointing at its RFC 9728 metadata. This is the canonical mechanism;
    // well-known guesses below are best-effort fallbacks for servers that
    // don't follow the spec.
    let www_authenticate_url =
        probe_resource_metadata_via_www_authenticate(client, server_url).await;

    // --- Step 1: RFC 9728 resource metadata ---
    // Use the URL from Step 0 if we have it; otherwise synthesize the
    // path-aware/origin well-known URL.
    let rfc9728_url = if let Some(u) = www_authenticate_url {
        u
    } else if path.is_empty() || path == "/" {
        format!("{origin}/.well-known/oauth-protected-resource")
    } else {
        format!("{origin}/.well-known/oauth-protected-resource{path}")
    };

    debug!("discover_as: trying RFC 9728 resource metadata at {rfc9728_url}");
    let as_url: Option<String> = match client.get(&rfc9728_url).send().await {
        Ok(resp) => {
            let status = resp.status();
            if status.is_success() {
                let meta: ResourceMetadata = resp.json().await.map_err(|e| {
                    OAuthError::DiscoveryFailed(format!(
                        "failed to parse RFC 9728 response from {rfc9728_url}: {e}"
                    ))
                })?;
                debug!(
                    "discover_as: RFC 9728 succeeded, AS URL = {:?}",
                    meta.authorization_servers.first()
                );
                Some(
                    meta.authorization_servers
                        .into_iter()
                        .next()
                        .ok_or_else(|| {
                            OAuthError::DiscoveryFailed(
                                "RFC 9728 authorization_servers array is empty".to_string(),
                            )
                        })?,
                )
            } else {
                // Any non-2xx (404, 401, 5xx, …) → speculative URL didn't
                // hit metadata. Skip and try the AS metadata fallback chain.
                debug!("discover_as: RFC 9728 returned {status}, falling back to AS metadata");
                None
            }
        }
        Err(e) => {
            // Connection error on a speculative URL → skip, try fallback.
            debug!("discover_as: RFC 9728 request failed for {rfc9728_url}: {e}, falling back");
            None
        }
    };

    // --- Steps 2 & 3: AS metadata + OIDC ---
    let as_base = as_url.as_deref().unwrap_or(server_url);
    let as_meta_urls = as_metadata_urls(as_base);

    let mut last_err: Option<String> = None;
    for url in &as_meta_urls {
        debug!("discover_as: trying AS metadata at {url}");
        match client.get(url).send().await {
            Ok(resp) => {
                let status = resp.status();
                if status.is_success() {
                    let meta: AsMetadata = resp.json().await.map_err(|e| {
                        OAuthError::DiscoveryFailed(format!(
                            "failed to parse AS metadata from {url}: {e}"
                        ))
                    })?;
                    debug!("discover_as: succeeded via {url}");
                    return Ok(meta);
                }
                debug!("discover_as: {url} returned {status}, trying next");
                last_err = Some(format!("{status} at {url}"));
            }
            Err(e) => {
                debug!("discover_as: AS metadata request failed for {url}: {e}, trying next");
                last_err = Some(format!("request failed for {url}: {e}"));
            }
        }
    }

    Err(OAuthError::DiscoveryFailed(format!(
        "no AS metadata found for {server_url} (last probe: {})",
        last_err.unwrap_or_else(|| "no probes executed".to_string())
    )))
}

// ---------------------------------------------------------------------------
// Task 2: DCR, auth URL builder, and token exchange
// ---------------------------------------------------------------------------

/// DCR request body (RFC 7591).
#[derive(Debug, Serialize)]
struct DcrRequest<'a> {
    client_name: &'a str,
    redirect_uris: Vec<&'a str>,
    grant_types: Vec<&'a str>,
    response_types: Vec<&'a str>,
    token_endpoint_auth_method: &'a str,
    application_type: &'a str,
}

/// Register this client via Dynamic Client Registration (RFC 7591).
///
/// POSTs to `registration_endpoint` with the standard native-app fields.
/// Returns the server-assigned `client_id` (and optional `client_secret`).
pub async fn register_client(
    client: &reqwest::Client,
    registration_endpoint: &str,
    redirect_uri: &str,
) -> Result<DcrResponse, OAuthError> {
    let body = DcrRequest {
        client_name: "Right Agent",
        redirect_uris: vec![redirect_uri],
        grant_types: vec!["authorization_code"],
        response_types: vec!["code"],
        token_endpoint_auth_method: "none",
        application_type: "native",
    };

    debug!("register_client: POST to {registration_endpoint}");
    let resp = client
        .post(registration_endpoint)
        .json(&body)
        .send()
        .await
        .map_err(|e| OAuthError::DcrFailed(format!("DCR request failed: {e}")))?;

    let status = resp.status();
    if !status.is_success() {
        let text = resp.text().await.unwrap_or_default();
        return Err(OAuthError::DcrFailed(format!(
            "DCR returned {status} from {registration_endpoint}: {text}"
        )));
    }

    let dcr: DcrResponse = resp.json().await.map_err(|e| {
        OAuthError::DcrFailed(format!(
            "failed to parse DCR response from {registration_endpoint}: {e}"
        ))
    })?;

    Ok(dcr)
}

/// Try DCR; fall back to a static `client_id` from `mcp.json` configuration.
///
/// Per D-08: if `registration_endpoint` is absent, use `static_client_id`.
/// If neither is available, return `OAuthError::MissingClientId`.
pub async fn register_client_or_fallback(
    client: &reqwest::Client,
    metadata: &AsMetadata,
    static_client_id: Option<&str>,
    redirect_uri: &str,
) -> Result<(String, Option<String>), OAuthError> {
    if let Some(reg_ep) = &metadata.registration_endpoint {
        let dcr = register_client(client, reg_ep, redirect_uri).await?;
        Ok((dcr.client_id, dcr.client_secret))
    } else if let Some(id) = static_client_id {
        debug!("register_client_or_fallback: using static client_id (no registration_endpoint)");
        Ok((id.to_string(), None))
    } else {
        Err(OAuthError::MissingClientId)
    }
}

/// Build the authorization URL with PKCE S256 parameters.
///
/// URL-encodes all parameter values. Appends `scope` only when `Some`.
pub fn build_auth_url(
    metadata: &AsMetadata,
    client_id: &str,
    redirect_uri: &str,
    state: &str,
    code_challenge: &str,
    scope: Option<&str>,
) -> String {
    let encode = |s: &str| {
        s.chars()
            .flat_map(|c| match c {
                'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '_' | '.' | '~' => {
                    vec![c as u8]
                }
                _ => format!("%{:02X}", c as u8).into_bytes(),
            })
            .map(|b| b as char)
            .collect::<String>()
    };

    let mut url = format!(
        "{}?response_type=code&client_id={}&redirect_uri={}&state={}&code_challenge={}&code_challenge_method=S256",
        metadata.authorization_endpoint,
        encode(client_id),
        encode(redirect_uri),
        encode(state),
        encode(code_challenge),
    );

    if let Some(s) = scope {
        url.push_str(&format!("&scope={}", encode(s)));
    }

    url
}

/// Exchange an authorization code for tokens (authorization_code grant with PKCE).
///
/// Sends a form-encoded POST to `token_endpoint`.
/// Includes `client_secret` in the form body when `Some`.
pub async fn exchange_token(
    client: &reqwest::Client,
    token_endpoint: &str,
    code: &str,
    redirect_uri: &str,
    client_id: &str,
    client_secret: Option<&str>,
    code_verifier: &str,
) -> Result<TokenResponse, OAuthError> {
    let mut params = vec![
        ("grant_type", "authorization_code"),
        ("code", code),
        ("redirect_uri", redirect_uri),
        ("client_id", client_id),
        ("code_verifier", code_verifier),
    ];

    let secret_owned;
    if let Some(s) = client_secret {
        secret_owned = s.to_string();
        params.push(("client_secret", &secret_owned));
    }

    debug!("exchange_token: POST to {token_endpoint}");
    let resp = client
        .post(token_endpoint)
        .form(&params)
        .send()
        .await
        .map_err(|e| {
            OAuthError::TokenExchangeFailed(format!("token exchange request failed: {e}"))
        })?;

    let status = resp.status();
    if !status.is_success() {
        let text = resp.text().await.unwrap_or_default();
        return Err(OAuthError::TokenExchangeFailed(format!(
            "token endpoint returned {status}: {text}"
        )));
    }

    let token_resp: TokenResponse = resp.json().await.map_err(|e| {
        OAuthError::TokenExchangeFailed(format!(
            "failed to parse token response from {token_endpoint}: {e}"
        ))
    })?;

    Ok(token_resp)
}

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

    /// Install ring as the rustls process-level crypto provider. Idempotent —
    /// safe to call from multiple tests in the same binary.
    fn setup_crypto() {
        // install_default returns Err(existing provider Arc) when already
        // installed by another test in the same binary — that's not a failure.
        let _ = rustls::crypto::ring::default_provider().install_default();
    }

    #[test]
    fn generate_pkce_verifier_is_43_chars() {
        let (verifier, _) = generate_pkce();
        assert_eq!(
            verifier.len(),
            43,
            "code_verifier should be 43 chars (32 bytes base64url-no-pad)"
        );
    }

    #[test]
    fn generate_pkce_challenge_is_43_chars() {
        let (_, challenge) = generate_pkce();
        assert_eq!(
            challenge.len(),
            43,
            "code_challenge should be 43 chars (SHA-256 base64url-no-pad)"
        );
    }

    #[test]
    fn generate_pkce_challenge_matches_s256_of_verifier() {
        let (verifier, challenge) = generate_pkce();
        let expected_challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
        assert_eq!(
            challenge, expected_challenge,
            "code_challenge must be S256(code_verifier)"
        );
    }

    #[test]
    fn generate_state_is_22_chars() {
        let state = generate_state();
        assert_eq!(
            state.len(),
            22,
            "state should be 22 chars (16 bytes base64url-no-pad)"
        );
    }

    #[test]
    fn verify_state_returns_true_for_matching() {
        let state = generate_state();
        assert!(
            verify_state(&state, &state),
            "matching states should return true"
        );
    }

    #[test]
    fn verify_state_returns_false_for_nonmatching() {
        let a = generate_state();
        let b = generate_state();
        // Extremely unlikely to collide (128-bit random), but handle it defensively
        if a != b {
            assert!(
                !verify_state(&a, &b),
                "different states should return false"
            );
        }
    }

    #[test]
    fn verify_state_returns_false_for_different_lengths() {
        let state = generate_state();
        let short = &state[..10];
        assert!(
            !verify_state(&state, short),
            "different-length states should return false"
        );
    }

    // Task 1: AS discovery URL construction tests
    #[test]
    fn discovery_urls_with_path_rfc9728_appends_path() {
        let urls = discovery_urls("https://mcp.example.com/mcp");
        // RFC 9728: path appended after .well-known/oauth-protected-resource
        assert_eq!(
            urls[0],
            "https://mcp.example.com/.well-known/oauth-protected-resource/mcp"
        );
    }

    #[test]
    fn discovery_urls_without_path_rfc9728_no_trailing_slash() {
        let urls = discovery_urls("https://api.example.com");
        assert_eq!(
            urls[0],
            "https://api.example.com/.well-known/oauth-protected-resource"
        );
    }

    #[test]
    fn discovery_urls_rfc8414_no_path() {
        let urls = discovery_urls("https://api.example.com");
        assert!(
            urls.contains(
                &"https://api.example.com/.well-known/oauth-authorization-server".to_string()
            ),
            "RFC 8414 no-path URL must be in discovery list"
        );
    }

    #[test]
    fn discovery_urls_oidc_no_path() {
        let urls = discovery_urls("https://api.example.com");
        assert!(
            urls.contains(&"https://api.example.com/.well-known/openid-configuration".to_string()),
            "OIDC no-path URL must be in discovery list"
        );
    }

    #[test]
    fn discovery_urls_rfc9728_order_is_first() {
        let urls = discovery_urls("https://mcp.example.com/mcp");
        // RFC 9728 must be first
        assert!(urls[0].contains("oauth-protected-resource"));
    }

    #[test]
    fn discovery_urls_all_contain_well_known() {
        let urls = discovery_urls("https://mcp.example.com/mcp");
        for url in &urls {
            assert!(
                url.contains(".well-known"),
                "All discovery URLs must contain .well-known: {url}"
            );
        }
    }

    // Task 1: discover_as integration tests using mock HTTP server
    #[tokio::test]
    async fn discover_as_rfc9728_success_returns_as_metadata() {
        setup_crypto();
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let as_meta_body = r#"{"authorization_endpoint":"https://auth.example.com/authorize","token_endpoint":"https://auth.example.com/token"}"#;

        // Spin up a minimal HTTP server
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();

        let as_meta_body_owned = as_meta_body.to_string();
        tokio::spawn(async move {
            // Step 0 probe: GET /mcp → 404 (no WWW-Authenticate header).
            // discover_as sends this first; a 404 means probe returns None.
            let (mut stream0, _) = listener.accept().await.unwrap();
            let mut buf0 = [0u8; 4096];
            let _ = stream0.read(&mut buf0).await;
            let resp0 = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
            let _ = stream0.write_all(resp0.as_bytes()).await;

            // Step 1: GET /.well-known/oauth-protected-resource/mcp → 200 with
            // ResourceMetadata pointing back to the same host for AS metadata.
            let (mut stream1, _) = listener.accept().await.unwrap();
            let mut buf1 = [0u8; 4096];
            let _ = stream1.read(&mut buf1).await;
            let rfc9728_body =
                format!(r#"{{"authorization_servers":["http://127.0.0.1:{port}"]}}"#);
            let resp1 = format!(
                "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                rfc9728_body.len(),
                rfc9728_body
            );
            let _ = stream1.write_all(resp1.as_bytes()).await;

            // Step 2+3: GET /.well-known/oauth-authorization-server → 200 with AS metadata.
            let (mut stream2, _) = listener.accept().await.unwrap();
            let mut buf2 = [0u8; 4096];
            let _ = stream2.read(&mut buf2).await;
            let resp2 = format!(
                "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                as_meta_body_owned.len(),
                as_meta_body_owned
            );
            let _ = stream2.write_all(resp2.as_bytes()).await;
        });

        let client = reqwest::Client::new();
        let server_url = format!("http://127.0.0.1:{port}/mcp");
        let result = discover_as(&client, &server_url).await;
        assert!(result.is_ok(), "discover_as should succeed: {result:?}");
        let meta = result.unwrap();
        assert_eq!(
            meta.authorization_endpoint,
            "https://auth.example.com/authorize"
        );
        assert_eq!(meta.token_endpoint, "https://auth.example.com/token");
    }

    #[tokio::test]
    async fn discover_as_5xx_skips_and_continues() {
        setup_crypto();
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;

        // RFC 9728 path-suffixed → 500 (must NOT abort the chain).
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-protected-resource/mcp"))
            .respond_with(ResponseTemplate::new(500))
            .mount(&server)
            .await;

        // Origin-only AS metadata → 200 (must be reached).
        let as_meta = serde_json::json!({
            "authorization_endpoint": "https://auth.example.com/authorize",
            "token_endpoint":         "https://auth.example.com/token"
        });
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-authorization-server"))
            .respond_with(ResponseTemplate::new(200).set_body_json(as_meta))
            .mount(&server)
            .await;

        // Step 0 probes GET /mcp → 404 (wiremock default), probe returns None.
        // Anything else 404s implicitly via wiremock's default.

        let client = reqwest::Client::new();
        let server_url = format!("{}/mcp", server.uri());
        let result = discover_as(&client, &server_url).await;

        let meta = result.expect("5xx on a speculative probe must not abort the chain");
        assert_eq!(meta.token_endpoint, "https://auth.example.com/token");
    }

    #[tokio::test]
    async fn discover_as_all_non_2xx_returns_discovery_failed() {
        setup_crypto();
        use wiremock::matchers::method;
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        // Step 0 gets 401 with no WWW-Authenticate header — probe falls through.
        // Well-known chain also gets 401 everywhere — chain has no successful URL.
        Mock::given(method("GET"))
            .respond_with(ResponseTemplate::new(401))
            .mount(&server)
            .await;

        let client = reqwest::Client::new();
        let server_url = format!("{}/mcp", server.uri());
        let result = discover_as(&client, &server_url).await;

        assert!(
            matches!(result, Err(OAuthError::DiscoveryFailed(_))),
            "all-401 chain must end in DiscoveryFailed: {result:?}"
        );
    }

    #[tokio::test]
    async fn discover_as_all_404_returns_discovery_failed() {
        setup_crypto();
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();

        // Serve 404 for all requests
        tokio::spawn(async move {
            for _ in 0..10 {
                if let Ok((mut stream, _)) = listener.accept().await {
                    let mut buf = [0u8; 4096];
                    let _ = stream.read(&mut buf).await;
                    let resp =
                        "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
                    let _ = stream.write_all(resp.as_bytes()).await;
                }
            }
        });

        let client = reqwest::Client::new();
        let server_url = format!("http://127.0.0.1:{port}");
        let result = discover_as(&client, &server_url).await;
        assert!(
            matches!(result, Err(OAuthError::DiscoveryFailed(_))),
            "all-404 should fail: {result:?}"
        );
    }

    // Task 2: build_auth_url tests (pure function, no HTTP needed)
    #[test]
    fn build_auth_url_contains_required_params() {
        let metadata = AsMetadata {
            authorization_endpoint: "https://auth.example.com/authorize".to_string(),
            token_endpoint: "https://auth.example.com/token".to_string(),
            registration_endpoint: None,
            code_challenge_methods_supported: Some(vec!["S256".to_string()]),
        };
        let url = build_auth_url(
            &metadata,
            "client123",
            "https://cb.example.com/callback",
            "state456",
            "challenge789",
            Some("read write"),
        );
        assert!(
            url.contains("response_type=code"),
            "must have response_type=code"
        );
        assert!(url.contains("client_id=client123"), "must have client_id");
        assert!(url.contains("redirect_uri="), "must have redirect_uri");
        assert!(url.contains("state=state456"), "must have state");
        assert!(
            url.contains("code_challenge=challenge789"),
            "must have code_challenge"
        );
        assert!(
            url.contains("code_challenge_method=S256"),
            "must have S256 method"
        );
        assert!(url.contains("scope="), "must have scope when provided");
    }

    #[test]
    fn build_auth_url_no_scope_when_none() {
        let metadata = AsMetadata {
            authorization_endpoint: "https://auth.example.com/authorize".to_string(),
            token_endpoint: "https://auth.example.com/token".to_string(),
            registration_endpoint: None,
            code_challenge_methods_supported: None,
        };
        let url = build_auth_url(
            &metadata,
            "client123",
            "https://cb.example.com/callback",
            "state456",
            "challenge789",
            None,
        );
        assert!(
            !url.contains("scope="),
            "scope param must be absent when None"
        );
    }

    // Task 2: register_client tests
    #[tokio::test]
    async fn register_client_posts_correct_body_and_returns_client_id() {
        setup_crypto();
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();

        tokio::spawn(async move {
            let (mut stream, _) = listener.accept().await.unwrap();
            let mut buf = [0u8; 8192];
            let n = stream.read(&mut buf).await.unwrap();
            let request = String::from_utf8_lossy(&buf[..n]);
            // Verify Right Agent client_name is in body
            assert!(
                request.contains("Right Agent"),
                "DCR body must contain Right Agent"
            );
            assert!(
                request.contains("authorization_code"),
                "DCR body must contain authorization_code"
            );

            let body = r#"{"client_id":"dcr-client-id-123","client_secret":null}"#;
            let resp = format!(
                "HTTP/1.1 201 Created\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                body.len(),
                body
            );
            let _ = stream.write_all(resp.as_bytes()).await;
        });

        let client = reqwest::Client::new();
        let reg_ep = format!("http://127.0.0.1:{port}/register");
        let result = register_client(&client, &reg_ep, "https://cb.example.com/callback").await;
        assert!(result.is_ok(), "register_client should succeed: {result:?}");
        assert_eq!(result.unwrap().client_id, "dcr-client-id-123");
    }

    #[tokio::test]
    async fn register_client_non_2xx_returns_dcr_failed() {
        setup_crypto();
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();

        tokio::spawn(async move {
            let (mut stream, _) = listener.accept().await.unwrap();
            let mut buf = [0u8; 4096];
            let _ = stream.read(&mut buf).await;
            let resp = "HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
            let _ = stream.write_all(resp.as_bytes()).await;
        });

        let client = reqwest::Client::new();
        let reg_ep = format!("http://127.0.0.1:{port}/register");
        let result = register_client(&client, &reg_ep, "https://cb.example.com/callback").await;
        assert!(
            matches!(result, Err(OAuthError::DcrFailed(_))),
            "non-2xx should be DcrFailed: {result:?}"
        );
    }

    #[test]
    fn register_client_or_fallback_uses_static_client_id_when_no_registration_endpoint() {
        // Synchronous test: when no registration_endpoint and static_client_id provided, returns static
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let client = reqwest::Client::new();
            let metadata = AsMetadata {
                authorization_endpoint: "https://auth.example.com/authorize".to_string(),
                token_endpoint: "https://auth.example.com/token".to_string(),
                registration_endpoint: None,
                code_challenge_methods_supported: None,
            };
            let result = register_client_or_fallback(
                &client,
                &metadata,
                Some("static-id-abc"),
                "https://cb.example.com/callback",
            )
            .await;
            assert!(result.is_ok(), "static fallback should succeed");
            let (client_id, secret) = result.unwrap();
            assert_eq!(client_id, "static-id-abc");
            assert!(secret.is_none());
        });
    }

    #[test]
    fn register_client_or_fallback_returns_missing_client_id_when_neither() {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let client = reqwest::Client::new();
            let metadata = AsMetadata {
                authorization_endpoint: "https://auth.example.com/authorize".to_string(),
                token_endpoint: "https://auth.example.com/token".to_string(),
                registration_endpoint: None,
                code_challenge_methods_supported: None,
            };
            let result = register_client_or_fallback(
                &client,
                &metadata,
                None,
                "https://cb.example.com/callback",
            )
            .await;
            assert!(
                matches!(result, Err(OAuthError::MissingClientId)),
                "no endpoint + no static id = MissingClientId"
            );
        });
    }

    // Task 2: exchange_token tests
    #[tokio::test]
    async fn exchange_token_posts_form_and_returns_token_response() {
        setup_crypto();
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();

        tokio::spawn(async move {
            let (mut stream, _) = listener.accept().await.unwrap();
            let mut buf = [0u8; 8192];
            let n = stream.read(&mut buf).await.unwrap();
            let request = String::from_utf8_lossy(&buf[..n]);
            // Verify form fields
            assert!(
                request.contains("grant_type=authorization_code"),
                "must have grant_type"
            );
            assert!(request.contains("code=auth-code-xyz"), "must have code");
            assert!(
                request.contains("code_verifier="),
                "must have code_verifier"
            );

            let body = r#"{"access_token":"tok-abc","token_type":"Bearer","expires_in":3600}"#;
            let resp = format!(
                "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                body.len(),
                body
            );
            let _ = stream.write_all(resp.as_bytes()).await;
        });

        let client = reqwest::Client::new();
        let token_ep = format!("http://127.0.0.1:{port}/token");
        let result = exchange_token(
            &client,
            &token_ep,
            "auth-code-xyz",
            "https://cb.example.com/callback",
            "client123",
            None,
            "verifier-string",
        )
        .await;
        assert!(result.is_ok(), "exchange_token should succeed: {result:?}");
        let tok = result.unwrap();
        assert_eq!(tok.access_token, "tok-abc");
        assert_eq!(tok.expires_in, Some(3600));
    }

    #[tokio::test]
    async fn exchange_token_non_2xx_returns_token_exchange_failed() {
        setup_crypto();
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();

        tokio::spawn(async move {
            let (mut stream, _) = listener.accept().await.unwrap();
            let mut buf = [0u8; 4096];
            let _ = stream.read(&mut buf).await;
            let body = r#"{"error":"invalid_grant"}"#;
            let resp = format!(
                "HTTP/1.1 400 Bad Request\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                body.len(),
                body
            );
            let _ = stream.write_all(resp.as_bytes()).await;
        });

        let client = reqwest::Client::new();
        let token_ep = format!("http://127.0.0.1:{port}/token");
        let result = exchange_token(
            &client,
            &token_ep,
            "bad-code",
            "https://cb.example.com/callback",
            "client123",
            None,
            "verifier",
        )
        .await;
        assert!(
            matches!(result, Err(OAuthError::TokenExchangeFailed(_))),
            "non-2xx should be TokenExchangeFailed: {result:?}"
        );
    }

    #[tokio::test]
    async fn discover_as_linear_pattern_uses_origin_well_known() {
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;

        // Step 0 probes GET /mcp → 404 (no mock matches; wiremock default),
        // probe returns None and the chain proceeds to the well-known fallback.

        // Path-suffixed RFC 9728 → 404 (Linear does not host metadata here)
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-protected-resource/mcp"))
            .respond_with(ResponseTemplate::new(404))
            .mount(&server)
            .await;

        // Path-suffixed RFC 8414 → 404
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-authorization-server/mcp"))
            .respond_with(ResponseTemplate::new(404))
            .mount(&server)
            .await;

        // Path-suffixed OIDC → 404
        Mock::given(method("GET"))
            .and(path("/.well-known/openid-configuration/mcp"))
            .respond_with(ResponseTemplate::new(404))
            .mount(&server)
            .await;

        // Resource-path-prefixed OIDC → 401 (this is the URL that aborted
        // the loop in production; mock Linear's actual response).
        Mock::given(method("GET"))
            .and(path("/mcp/.well-known/openid-configuration"))
            .respond_with(ResponseTemplate::new(401))
            .mount(&server)
            .await;

        // ORIGIN-ONLY RFC 8414 → 200 with valid AS metadata. Linear hosts
        // metadata here. Current code never probes this URL when the MCP
        // server URL has a non-empty path; Task 2 makes it probe this URL,
        // and Task 3 makes the chain reach this URL despite the 401 above.
        let as_meta = serde_json::json!({
            "authorization_endpoint": "https://auth.linear.example/authorize",
            "token_endpoint":         "https://auth.linear.example/token"
        });
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-authorization-server"))
            .respond_with(ResponseTemplate::new(200).set_body_json(as_meta))
            .mount(&server)
            .await;

        let client = reqwest::Client::new();
        let server_url = format!("{}/mcp", server.uri());
        let result = discover_as(&client, &server_url).await;

        let meta = result.expect("discover_as must succeed for Linear-style server");
        assert_eq!(
            meta.authorization_endpoint,
            "https://auth.linear.example/authorize"
        );
        assert_eq!(meta.token_endpoint, "https://auth.linear.example/token");
    }

    #[test]
    fn as_metadata_urls_path_bearing_includes_origin_only_fallback() {
        let urls = as_metadata_urls("https://mcp.linear.app/mcp");
        assert_eq!(urls.len(), 5, "expected 5 URLs, got {urls:?}");
        // Path-aware variants come first (RFC 8414 §3.1).
        assert_eq!(
            urls[0],
            "https://mcp.linear.app/.well-known/oauth-authorization-server/mcp"
        );
        assert_eq!(
            urls[1],
            "https://mcp.linear.app/.well-known/openid-configuration/mcp"
        );
        assert_eq!(
            urls[2],
            "https://mcp.linear.app/mcp/.well-known/openid-configuration"
        );
        // Origin-only fallbacks come after (real-world: Linear, …).
        assert_eq!(
            urls[3],
            "https://mcp.linear.app/.well-known/oauth-authorization-server"
        );
        assert_eq!(
            urls[4],
            "https://mcp.linear.app/.well-known/openid-configuration"
        );
    }

    #[test]
    fn as_metadata_urls_no_path_unchanged() {
        // Origin-only AS URL should still produce the same two URLs as before.
        let urls = as_metadata_urls("https://auth.example.com");
        assert_eq!(urls.len(), 2);
        assert_eq!(
            urls[0],
            "https://auth.example.com/.well-known/oauth-authorization-server"
        );
        assert_eq!(
            urls[1],
            "https://auth.example.com/.well-known/openid-configuration"
        );
    }

    #[test]
    fn parse_www_authenticate_quoted_resource_metadata() {
        let h = r#"Bearer realm="Linear", resource_metadata="https://mcp.linear.app/.well-known/oauth-protected-resource""#;
        assert_eq!(
            parse_www_authenticate_resource_metadata(h),
            Some("https://mcp.linear.app/.well-known/oauth-protected-resource".to_string())
        );
    }

    #[test]
    fn parse_www_authenticate_unquoted_resource_metadata() {
        let h = "Bearer resource_metadata=https://api.example.com/.well-known/x";
        assert_eq!(
            parse_www_authenticate_resource_metadata(h),
            Some("https://api.example.com/.well-known/x".to_string())
        );
    }

    #[test]
    fn parse_www_authenticate_unquoted_with_trailing_param() {
        let h = "Bearer resource_metadata=https://api.example.com/x, error=invalid_token";
        assert_eq!(
            parse_www_authenticate_resource_metadata(h),
            Some("https://api.example.com/x".to_string())
        );
    }

    #[test]
    fn parse_www_authenticate_no_resource_metadata_returns_none() {
        assert_eq!(
            parse_www_authenticate_resource_metadata(r#"Bearer realm="x""#),
            None
        );
        assert_eq!(
            parse_www_authenticate_resource_metadata("Basic realm=x"),
            None
        );
        assert_eq!(parse_www_authenticate_resource_metadata(""), None);
    }

    #[test]
    fn parse_www_authenticate_empty_quoted_value_returns_none() {
        assert_eq!(
            parse_www_authenticate_resource_metadata(r#"Bearer resource_metadata="""#),
            None,
            "empty quoted value is malformed and must be None"
        );
    }

    #[tokio::test]
    async fn discover_as_uses_www_authenticate_resource_metadata() {
        setup_crypto();
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        let metadata_path = "/.well-known/oauth-protected-resource";

        // Step 0: unauthenticated GET /mcp returns 401 with WWW-Authenticate
        // pointing to the metadata URL on this same host.
        let www_authenticate = format!(
            r#"Bearer realm="MCP", resource_metadata="{}{}""#,
            server.uri(),
            metadata_path
        );
        Mock::given(method("GET"))
            .and(path("/mcp"))
            .respond_with(
                ResponseTemplate::new(401)
                    .insert_header("WWW-Authenticate", www_authenticate.as_str()),
            )
            .expect(1)
            .mount(&server)
            .await;

        // Step 1: RFC 9728 resource metadata at the WWW-Authenticate-supplied URL
        // points at the AS URL, also on this same host.
        let resource_meta = serde_json::json!({
            "authorization_servers": [server.uri()]
        });
        Mock::given(method("GET"))
            .and(path(metadata_path))
            .respond_with(ResponseTemplate::new(200).set_body_json(resource_meta))
            .expect(1)
            .mount(&server)
            .await;

        // Step 2: AS metadata at the origin-only RFC 8414 URL of the AS host.
        let as_meta = serde_json::json!({
            "authorization_endpoint": "https://auth.example.com/authorize",
            "token_endpoint":         "https://auth.example.com/token"
        });
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-authorization-server"))
            .respond_with(ResponseTemplate::new(200).set_body_json(as_meta))
            .expect(1)
            .mount(&server)
            .await;

        let client = reqwest::Client::new();
        let server_url = format!("{}/mcp", server.uri());
        let result = discover_as(&client, &server_url).await;

        let meta = result.expect("WWW-Authenticate-driven discovery must succeed");
        assert_eq!(
            meta.authorization_endpoint,
            "https://auth.example.com/authorize"
        );
        assert_eq!(meta.token_endpoint, "https://auth.example.com/token");
    }
}
