//! AWS Signature Version 4 (SigV4) request signing for URL targets.
//!
//! Signs an outbound request with the AWS SigV4 algorithm using the
//! `aws-sigv4` SDK crate. Credentials resolve from the environment per
//! the configured source. Signatures are time-bound, so we sign every
//! request fresh; the transport calls [`sign_request`] just before send.
//!
//! See `docs/guides/auth/auth.md` for the YAML shape and
//! `docs/adr/0031-sigv4-auth.md` for the design rationale.

use std::time::{Duration, SystemTime};

use aws_credential_types::Credentials;
use aws_sigv4::http_request::{sign, SignableBody, SignableRequest, SigningSettings};
use aws_sigv4::sign::v4;
use serde::{Deserialize, Serialize};

use super::AuthError;

/// AWS SigV4 configuration for a URL-target server.
///
/// Mirrors the AWS Rust SDK credential chain. Most CI environments
/// resolve credentials through environment variables (`AWS_ACCESS_KEY_ID`,
/// etc.); workstations typically use a named profile from `~/.aws/credentials`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SigV4Config {
    /// AWS region the signed request is destined for.
    pub region: String,
    /// AWS service identifier (e.g. `execute-api` for API Gateway,
    /// `lambda` for Function URLs).
    pub service: String,
    /// Credential source.
    pub credentials: SigV4Credentials,
    /// Optional override for the signature expiry window. Defaults to
    /// 15 minutes which matches the SigV4 spec; rarely needs changing.
    #[serde(
        default,
        with = "humantime_serde",
        skip_serializing_if = "Option::is_none"
    )]
    pub signature_lifetime: Option<Duration>,
}

mod humantime_serde {
    use std::time::Duration;

    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    pub fn serialize<S: Serializer>(
        value: &Option<Duration>,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        match value {
            Some(d) => format!("{}s", d.as_secs()).serialize(serializer),
            None => serializer.serialize_none(),
        }
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(
        deserializer: D,
    ) -> Result<Option<Duration>, D::Error> {
        let s: Option<String> = Option::deserialize(deserializer)?;
        match s {
            None => Ok(None),
            Some(raw) => {
                let trimmed = raw.trim().trim_end_matches('s');
                trimmed
                    .parse::<u64>()
                    .map(|secs| Some(Duration::from_secs(secs)))
                    .map_err(serde::de::Error::custom)
            }
        }
    }
}

/// Where to get the IAM credentials from.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "source")]
pub enum SigV4Credentials {
    /// Walk the standard AWS SDK credential chain: env vars, named
    /// profile, IMDS, EC2 instance role, etc. This is the right default
    /// for nearly every deployment.
    FromEnvironment,
    /// Use a named profile from `~/.aws/credentials` or `~/.aws/config`.
    Profile {
        /// Profile name.
        name: String,
    },
    /// Read credentials from explicit environment variable names. Useful
    /// when CI exports them under non-standard names.
    Explicit {
        /// Env var holding the access key ID.
        access_key_id_env: String,
        /// Env var holding the secret access key.
        secret_access_key_env: String,
        /// Optional env var holding a session token (STS).
        #[serde(default, skip_serializing_if = "Option::is_none")]
        session_token_env: Option<String>,
    },
}

/// Standard AWS env var names the default chain reads.
const ENV_ACCESS_KEY: &str = "AWS_ACCESS_KEY_ID";
const ENV_SECRET_KEY: &str = "AWS_SECRET_ACCESS_KEY";
const ENV_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN";

impl SigV4Config {
    /// Validate the config without contacting AWS. Returns an error if
    /// `region` or `service` is empty.
    pub fn validate(&self) -> Result<(), AuthError> {
        if self.region.trim().is_empty() {
            return Err(AuthError::InvalidConfig {
                reason: "sigv4.region must not be empty".into(),
            });
        }
        if self.service.trim().is_empty() {
            return Err(AuthError::InvalidConfig {
                reason: "sigv4.service must not be empty".into(),
            });
        }
        Ok(())
    }
}

/// One resolved SigV4 header to attach to the outbound request.
///
/// Always includes `Authorization` and `X-Amz-Date`; includes
/// `X-Amz-Security-Token` when the credentials carry a session token.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SignedHeader {
    /// Header name (e.g. `authorization`).
    pub name: String,
    /// Header value.
    pub value: String,
}

/// Sign an outbound request with SigV4 and return the headers to attach.
///
/// `body` is the request body the signature commits to (empty for a GET).
/// Credentials resolve from the environment per [`SigV4Config::credentials`].
/// The signature is time-bound, so callers sign per request.
pub fn sign_request(
    method: &str,
    url: &str,
    body: &[u8],
    config: &SigV4Config,
) -> Result<Vec<SignedHeader>, AuthError> {
    let (creds, region) = resolve_credentials(&config.credentials)?;
    // A named profile can supply the region when the YAML/CLI did not. The
    // config region still wins when both are set, so we only fill an empty
    // one. Validation runs after the fill so a profile-only region passes.
    let mut effective = config.clone();
    if effective.region.trim().is_empty() {
        if let Some(region) = region {
            effective.region = region;
        }
    }
    effective.validate()?;
    sign_with_credentials(method, url, body, &effective, &creds, SystemTime::now())
}

/// Signing core with explicit credentials and timestamp. Exposed to the
/// test module so a frozen clock yields a deterministic signature.
fn sign_with_credentials(
    method: &str,
    url: &str,
    body: &[u8],
    config: &SigV4Config,
    creds: &Credentials,
    time: SystemTime,
) -> Result<Vec<SignedHeader>, AuthError> {
    let identity = creds.clone().into();
    let params = v4::SigningParams::builder()
        .identity(&identity)
        .region(&config.region)
        .name(&config.service)
        .time(time)
        .settings(SigningSettings::default())
        .build()
        .map_err(|e| crypto_error(&format!("invalid signing params: {e}")))?
        .into();

    let signable = SignableRequest::new(method, url, std::iter::empty(), SignableBody::Bytes(body))
        .map_err(|e| crypto_error(&format!("could not build signable request: {e}")))?;

    let (instructions, _signature) = sign(signable, &params)
        .map_err(|e| crypto_error(&format!("signing failed: {e}")))?
        .into_parts();

    let (headers, _params) = instructions.into_parts();
    Ok(headers
        .into_iter()
        .map(|h| SignedHeader {
            name: h.name().to_string(),
            value: h.value().to_string(),
        })
        .collect())
}

/// Resolve IAM credentials and an optional region from the configured
/// source. The default and explicit sources read the environment (no
/// region); the profile source reads `~/.aws/credentials` (and
/// `~/.aws/config` for the region) via [`super::sigv4_profile`].
fn resolve_credentials(
    source: &SigV4Credentials,
) -> Result<(Credentials, Option<String>), AuthError> {
    match source {
        SigV4Credentials::FromEnvironment => Ok((
            credentials_from_env(ENV_ACCESS_KEY, ENV_SECRET_KEY, Some(ENV_SESSION_TOKEN))?,
            None,
        )),
        SigV4Credentials::Explicit {
            access_key_id_env,
            secret_access_key_env,
            session_token_env,
        } => Ok((
            credentials_from_env(
                access_key_id_env,
                secret_access_key_env,
                session_token_env.as_deref(),
            )?,
            None,
        )),
        SigV4Credentials::Profile { name } => {
            let resolved = super::sigv4_profile::resolve_profile(name)?;
            let creds = Credentials::new(
                resolved.access_key_id,
                resolved.secret_access_key,
                resolved.session_token,
                None,
                "mcptest-sigv4-profile",
            );
            Ok((creds, resolved.region))
        }
    }
}

/// Read credentials from the named env vars. The session token var is
/// optional; when present but empty it is treated as absent.
fn credentials_from_env(
    access_key_var: &str,
    secret_key_var: &str,
    session_token_var: Option<&str>,
) -> Result<Credentials, AuthError> {
    let access_key = require_env(access_key_var)?;
    let secret_key = require_env(secret_key_var)?;
    let session_token = session_token_var
        .and_then(|name| std::env::var(name).ok())
        .filter(|v| !v.is_empty());
    Ok(Credentials::new(
        access_key,
        secret_key,
        session_token,
        None,
        "mcptest-sigv4-env",
    ))
}

fn require_env(name: &str) -> Result<String, AuthError> {
    std::env::var(name).map_err(|_| AuthError::MissingEnvVar {
        name: name.to_string(),
    })
}

fn crypto_error(reason: &str) -> AuthError {
    AuthError::Crypto {
        scheme: "aws_sigv4",
        reason: reason.to_string(),
    }
}

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

    fn fixed_credentials() -> Credentials {
        // AWS SigV4 test-suite example credentials. Not real secrets.
        Credentials::new(
            "AKIDEXAMPLE",
            "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
            None,
            None,
            "test",
        )
    }

    fn example_config() -> SigV4Config {
        SigV4Config {
            region: "us-east-1".into(),
            service: "execute-api".into(),
            credentials: SigV4Credentials::FromEnvironment,
            signature_lifetime: None,
        }
    }

    #[test]
    fn validate_rejects_empty_region() {
        let cfg = SigV4Config {
            region: String::new(),
            ..example_config()
        };
        assert!(matches!(
            cfg.validate(),
            Err(AuthError::InvalidConfig { .. })
        ));
    }

    #[test]
    fn validate_rejects_empty_service() {
        let cfg = SigV4Config {
            service: String::new(),
            ..example_config()
        };
        assert!(matches!(
            cfg.validate(),
            Err(AuthError::InvalidConfig { .. })
        ));
    }

    #[test]
    fn validate_accepts_well_formed_config() {
        assert!(example_config().validate().is_ok());
    }

    #[test]
    fn credentials_round_trip_through_serde() {
        let cfg = SigV4Config {
            region: "eu-west-1".into(),
            service: "lambda".into(),
            credentials: SigV4Credentials::Explicit {
                access_key_id_env: "AWS_ACCESS_KEY_ID".into(),
                secret_access_key_env: "AWS_SECRET_ACCESS_KEY".into(),
                session_token_env: Some("AWS_SESSION_TOKEN".into()),
            },
            signature_lifetime: None,
        };
        let yaml = serde_yaml::to_string(&cfg).unwrap();
        let decoded: SigV4Config = serde_yaml::from_str(&yaml).unwrap();
        assert_eq!(cfg, decoded);
    }

    #[test]
    fn signs_with_well_formed_authorization_header() {
        let frozen = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
        let headers = sign_with_credentials(
            "GET",
            "https://abc.execute-api.us-east-1.amazonaws.com/prod/mcp",
            &[],
            &example_config(),
            &fixed_credentials(),
            frozen,
        )
        .expect("signing should succeed");

        let auth = headers
            .iter()
            .find(|h| h.name.eq_ignore_ascii_case("authorization"))
            .expect("Authorization header present");
        assert!(auth.value.starts_with("AWS4-HMAC-SHA256 "));
        assert!(auth.value.contains("Credential=AKIDEXAMPLE/"));
        assert!(auth.value.contains("/us-east-1/execute-api/aws4_request"));
        assert!(auth.value.contains("SignedHeaders="));
        assert!(auth.value.contains("Signature="));

        assert!(headers
            .iter()
            .any(|h| h.name.eq_ignore_ascii_case("x-amz-date")));
    }

    #[test]
    fn signature_is_deterministic_for_a_frozen_clock() {
        let frozen = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
        let url = "https://abc.execute-api.us-east-1.amazonaws.com/prod/mcp";
        let first = sign_with_credentials(
            "POST",
            url,
            b"{}",
            &example_config(),
            &fixed_credentials(),
            frozen,
        )
        .unwrap();
        let second = sign_with_credentials(
            "POST",
            url,
            b"{}",
            &example_config(),
            &fixed_credentials(),
            frozen,
        )
        .unwrap();
        assert_eq!(first, second, "same inputs must yield the same signature");
    }

    #[test]
    fn session_token_credentials_emit_security_token_header() {
        let frozen = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000);
        let creds = Credentials::new(
            "AKIDEXAMPLE",
            "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
            Some("session-token-xyz".to_string()),
            None,
            "test",
        );
        let headers = sign_with_credentials(
            "GET",
            "https://abc.execute-api.us-east-1.amazonaws.com/prod/mcp",
            &[],
            &example_config(),
            &creds,
            frozen,
        )
        .unwrap();
        assert!(headers
            .iter()
            .any(|h| h.name.eq_ignore_ascii_case("x-amz-security-token")));
    }

    #[test]
    fn profile_source_resolves_through_the_profile_module() {
        let _lock = crate::auth::aws_env_lock();
        // Point the credentials file env override at a temp file so the
        // profile source reads real keys without touching the user's
        // `~/.aws`. This proves the dispatch reaches the profile resolver.
        let dir = tempfile::tempdir().unwrap();
        let creds_path = dir.path().join("credentials");
        std::fs::write(
            &creds_path,
            "[prod]\naws_access_key_id = AKIDEXAMPLE\naws_secret_access_key = not-a-real-secret\nregion = us-west-2\n",
        )
        .unwrap();
        std::env::set_var("AWS_SHARED_CREDENTIALS_FILE", &creds_path);
        std::env::set_var("AWS_CONFIG_FILE", dir.path().join("config"));
        let (creds, region) = resolve_credentials(&SigV4Credentials::Profile {
            name: "prod".into(),
        })
        .unwrap();
        std::env::remove_var("AWS_SHARED_CREDENTIALS_FILE");
        std::env::remove_var("AWS_CONFIG_FILE");
        assert_eq!(creds.access_key_id(), "AKIDEXAMPLE");
        assert_eq!(region.as_deref(), Some("us-west-2"));
    }

    #[test]
    fn profile_region_fills_empty_config_region() {
        let _lock = crate::auth::aws_env_lock();
        let dir = tempfile::tempdir().unwrap();
        let creds_path = dir.path().join("credentials");
        std::fs::write(
            &creds_path,
            "[prod]\naws_access_key_id = AKIDEXAMPLE\naws_secret_access_key = wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY\nregion = us-east-1\n",
        )
        .unwrap();
        std::env::set_var("AWS_SHARED_CREDENTIALS_FILE", &creds_path);
        std::env::set_var("AWS_CONFIG_FILE", dir.path().join("config"));
        let cfg = SigV4Config {
            region: String::new(),
            service: "execute-api".into(),
            credentials: SigV4Credentials::Profile {
                name: "prod".into(),
            },
            signature_lifetime: None,
        };
        let headers = sign_request(
            "GET",
            "https://abc.execute-api.us-east-1.amazonaws.com/prod/mcp",
            &[],
            &cfg,
        );
        std::env::remove_var("AWS_SHARED_CREDENTIALS_FILE");
        std::env::remove_var("AWS_CONFIG_FILE");
        let headers = headers.expect("profile region should satisfy validation");
        let auth = headers
            .iter()
            .find(|h| h.name.eq_ignore_ascii_case("authorization"))
            .expect("Authorization header present");
        // The profile-supplied region appears in the credential scope.
        assert!(auth.value.contains("/us-east-1/execute-api/aws4_request"));
    }
}
