use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use jsonwebtoken::{jwk::JwkSet, DecodingKey};

pub(super) type JwksFetchResult = Result<Vec<(String, DecodingKey)>, JwksError>;
pub(super) type JwksFetchFuture = Pin<Box<dyn Future<Output = JwksFetchResult> + Send>>;
pub(super) type JwksProviderRef = Arc<dyn JwksProvider>;

/// Abstraction for loading JWKS keys for initial load and periodic refresh.
pub(super) trait JwksProvider: Send + Sync {
    fn fetch_jwks(&self, client: reqwest::Client, uri: String) -> JwksFetchFuture;
}

/// Production provider that performs real JWKS HTTP fetches.
pub(super) struct RealJwksProvider;

impl JwksProvider for RealJwksProvider {
    fn fetch_jwks(&self, client: reqwest::Client, uri: String) -> JwksFetchFuture {
        Box::pin(async move { fetch_jwks(&client, &uri).await })
    }
}

#[cfg(test)]
pub(super) mod mock {
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;

    use super::{JwksFetchFuture, JwksFetchResult, JwksProvider};

    pub type MockJwksResponse = Arc<dyn Fn() -> JwksFetchResult + Send + Sync>;

    /// Deterministic in-memory JWKS provider for unit tests.
    ///
    /// Responses are returned in order and, once exhausted, the final response is reused
    /// for subsequent calls (useful for background refresh loops).
    pub struct MockJwksProvider {
        responses: Vec<MockJwksResponse>,
        call_count: AtomicUsize,
    }

    impl MockJwksProvider {
        pub fn new(responses: Vec<MockJwksResponse>) -> Self {
            assert!(
                !responses.is_empty(),
                "MockJwksProvider requires at least one response"
            );
            Self {
                responses,
                call_count: AtomicUsize::new(0),
            }
        }
    }

    impl JwksProvider for MockJwksProvider {
        fn fetch_jwks(&self, _client: reqwest::Client, _uri: String) -> JwksFetchFuture {
            let idx = self.call_count.fetch_add(1, Ordering::SeqCst);
            let response_factory = self
                .responses
                .get(idx)
                .unwrap_or_else(|| self.responses.last().unwrap())
                .clone();
            Box::pin(async move { (response_factory)() })
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub(super) enum JwksError {
    #[error("JWKS request failed: {0}")]
    RequestFailed(#[source] reqwest::Error),

    #[error("JWKS endpoint returned a retryable error status: {0}")]
    RetryableStatus(reqwest::StatusCode),

    #[error("JWKS endpoint returned an error status: {0}")]
    ErrorStatus(reqwest::StatusCode),

    #[error("invalid JWKS response: {0}")]
    InvalidResponse(#[source] reqwest::Error),

    #[error("JWK is missing required 'kid' field")]
    MissingKid,

    #[error("invalid JWK key material: {0}")]
    InvalidKey(#[from] jsonwebtoken::errors::Error),

    #[error("JWKS returned an empty key set")]
    EmptyKeySet,

    #[error("JWKS URI has a disallowed scheme: {0}")]
    DisallowedUriScheme(String),
}

fn validate_jwks_uri(uri: &str) -> Result<(), JwksError> {
    let parsed = reqwest::Url::parse(uri)
        .map_err(|err| JwksError::DisallowedUriScheme(format!("{uri} ({err})")))?;
    match parsed.scheme() {
        "https" => Ok(()),
        "http" if matches!(parsed.host_str(), Some("localhost" | "127.0.0.1")) => Ok(()),
        _ => Err(JwksError::DisallowedUriScheme(uri.to_owned())),
    }
}

/// TODO: Wrap with exponential backoff and jitter to
/// handle transient failures more gracefully instead of failing immediately.
pub(super) async fn fetch_jwks(
    client: &reqwest::Client,
    uri: &str,
) -> Result<Vec<(String, DecodingKey)>, JwksError> {
    validate_jwks_uri(uri)?;

    let response = client
        .get(uri)
        .send()
        .await
        .map_err(JwksError::RequestFailed)?;

    let status = response.status();
    if !status.is_success() {
        return Err(match status.as_u16() {
            429 | 502 | 503 | 504 => JwksError::RetryableStatus(status),
            _ => JwksError::ErrorStatus(status),
        });
    }

    let jwk_set = response
        .json::<JwkSet>()
        .await
        .map_err(JwksError::InvalidResponse)?;

    if jwk_set.keys.is_empty() {
        return Err(JwksError::EmptyKeySet);
    }

    let mut keys = Vec::new();
    for jwk in &jwk_set.keys {
        let kid = jwk.common.key_id.clone().ok_or(JwksError::MissingKid)?;
        let key = DecodingKey::from_jwk(jwk)?;
        keys.push((kid, key));
    }
    Ok(keys)
}

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

    /// Key ID used in test JWKS fixtures to identify the test RSA key.
    const TEST_KID: &str = "test-key-1";
    /// RSA public key modulus (base64url) paired with TEST_RSA_E to form a valid JWK.
    const TEST_RSA_N: &str = "rguts8kRMaf55JSXIfp0Nt1GYOEZeCaAwGiegSPZuKOh7OqzjhGCBayaNmoPEB86gRGKzpW9cdsSHAholAsQXiQYrMNZMgcDLT3TE_4n5KeL3zEsFL-lsV9QndM2chpZ3_U_jzreOMz_U0JbBS6WsNU2HYwsFW-vU69fTVqChKf4SfbM9rfV1I7zLZxEpYKyVIoSMXydTE90TuoPtXMtOLgRuBjp5En_szHqY3nkF-VbB9Fq1mhu1LsnuynbxTWErwWf43rqGt9BFyvouwyzdfOD9H69hxw5i09De9dXLAU7nl5C9AsOs_67DY4VFznMWRQ6AoFEhd-j5NzopVa3nQ";
    /// RSA public key exponent (base64url) — 65537, the standard public exponent.
    const TEST_RSA_E: &str = "AQAB";

    fn client() -> reqwest::Client {
        reqwest::Client::new()
    }

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

        let server = MockServer::start().await;
        let jwks = serde_json::json!({
            "keys": [{
                "kty": "RSA",
                "kid": TEST_KID,
                "use": "sig",
                "alg": "RS256",
                "n": TEST_RSA_N,
                "e": TEST_RSA_E
            }]
        });
        Mock::given(method("GET"))
            .and(path("/.well-known/jwks.json"))
            .respond_with(ResponseTemplate::new(200).set_body_json(&jwks))
            .mount(&server)
            .await;

        let uri = format!("{}/.well-known/jwks.json", server.uri());
        let keys = fetch_jwks(&client(), &uri).await.unwrap();

        assert_eq!(keys.len(), 1);
        assert_eq!(keys[0].0, TEST_KID);
    }

    #[tokio::test]
    /// Verifies fetch_jwks fails entirely when any key in the set is malformed,
    /// even if other keys are valid.
    async fn test_fetch_jwks_fails_fast_on_bad_key() {
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        let jwks = serde_json::json!({
            "keys": [
                {
                    "kty": "RSA",
                    "kid": TEST_KID,
                    "use": "sig",
                    "alg": "RS256",
                    "n": TEST_RSA_N,
                    "e": TEST_RSA_E
                },
                {
                    "kty": "RSA",
                    "kid": "bad-key",
                    "use": "sig",
                    "alg": "RS256",
                    "n": "!!!not-valid-base64url!!!",
                    "e": "AQAB"
                }
            ]
        });
        Mock::given(method("GET"))
            .and(path("/.well-known/jwks.json"))
            .respond_with(ResponseTemplate::new(200).set_body_json(&jwks))
            .mount(&server)
            .await;

        let uri = format!("{}/.well-known/jwks.json", server.uri());
        assert!(matches!(
            fetch_jwks(&client(), &uri).await.unwrap_err(),
            JwksError::InvalidKey(_)
        ));
    }

    #[tokio::test]
    async fn test_fetch_jwks_rejects_disallowed_uri_scheme() {
        assert!(matches!(
            fetch_jwks(&client(), "http://auth.example.com/.well-known/jwks.json")
                .await
                .unwrap_err(),
            JwksError::DisallowedUriScheme(_)
        ));
    }

    #[tokio::test]
    async fn test_fetch_jwks_errors_on_unreachable_host() {
        assert!(matches!(
            fetch_jwks(&client(), "http://127.0.0.1:1")
                .await
                .unwrap_err(),
            JwksError::RequestFailed(_)
        ));
    }

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

        let server = MockServer::start().await;
        let jwks = serde_json::json!({
            "keys": [{
                "kty": "RSA",
                "use": "sig",
                "alg": "RS256",
                "n": TEST_RSA_N,
                "e": TEST_RSA_E
            }]
        });
        Mock::given(method("GET"))
            .and(path("/.well-known/jwks.json"))
            .respond_with(ResponseTemplate::new(200).set_body_json(&jwks))
            .mount(&server)
            .await;

        let uri = format!("{}/.well-known/jwks.json", server.uri());
        assert!(matches!(
            fetch_jwks(&client(), &uri).await.unwrap_err(),
            JwksError::MissingKid
        ));
    }

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

        let server = MockServer::start().await;
        Mock::given(method("GET"))
            .and(path("/.well-known/jwks.json"))
            .respond_with(
                ResponseTemplate::new(200).set_body_json(&serde_json::json!({"keys": []})),
            )
            .mount(&server)
            .await;

        let uri = format!("{}/.well-known/jwks.json", server.uri());
        assert!(matches!(
            fetch_jwks(&client(), &uri).await.unwrap_err(),
            JwksError::EmptyKeySet
        ));
    }

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

        let server = MockServer::start().await;
        Mock::given(method("GET"))
            .and(path("/.well-known/jwks.json"))
            .respond_with(ResponseTemplate::new(200).set_body_string("not json"))
            .mount(&server)
            .await;

        let uri = format!("{}/.well-known/jwks.json", server.uri());
        assert!(matches!(
            fetch_jwks(&client(), &uri).await.unwrap_err(),
            JwksError::InvalidResponse(_)
        ));
    }

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

        let server = MockServer::start().await;
        Mock::given(method("GET"))
            .and(path("/.well-known/jwks.json"))
            .respond_with(ResponseTemplate::new(500))
            .mount(&server)
            .await;

        let uri = format!("{}/.well-known/jwks.json", server.uri());
        assert!(matches!(
            fetch_jwks(&client(), &uri).await.unwrap_err(),
            JwksError::ErrorStatus(_)
        ));
    }

    #[tokio::test]
    /// Verifies that a 500 response body containing valid JWKS JSON is still rejected.
    async fn test_fetch_jwks_error_status_ignores_valid_jwks_body() {
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        let jwks = serde_json::json!({
            "keys": [{
                "kty": "RSA",
                "kid": TEST_KID,
                "use": "sig",
                "alg": "RS256",
                "n": TEST_RSA_N,
                "e": TEST_RSA_E
            }]
        });
        Mock::given(method("GET"))
            .and(path("/.well-known/jwks.json"))
            .respond_with(ResponseTemplate::new(500).set_body_json(&jwks))
            .mount(&server)
            .await;

        let uri = format!("{}/.well-known/jwks.json", server.uri());
        assert!(matches!(
            fetch_jwks(&client(), &uri).await.unwrap_err(),
            JwksError::ErrorStatus(_)
        ));
    }

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

        // u16 suffix on the first element so Rust infers the array as [u16; 4]
        // rather than defaulting to i32, which ResponseTemplate::new won't accept.
        for status in [429u16, 502, 503, 504] {
            let server = MockServer::start().await;
            Mock::given(method("GET"))
                .and(path("/.well-known/jwks.json"))
                .respond_with(ResponseTemplate::new(status))
                .mount(&server)
                .await;

            let uri = format!("{}/.well-known/jwks.json", server.uri());
            assert!(
                matches!(
                    fetch_jwks(&client(), &uri).await.unwrap_err(),
                    JwksError::RetryableStatus(_)
                ),
                "expected RetryableStatus for HTTP {status}"
            );
        }
    }

    #[test]
    fn test_validate_jwks_uri_accepts_https() {
        assert!(validate_jwks_uri("https://auth.example.com/.well-known/jwks.json").is_ok());
    }

    #[test]
    fn test_validate_jwks_uri_accepts_http_localhost() {
        assert!(validate_jwks_uri("http://localhost/.well-known/jwks.json").is_ok());
    }

    #[test]
    fn test_validate_jwks_uri_accepts_http_127_0_0_1() {
        assert!(validate_jwks_uri("http://127.0.0.1/.well-known/jwks.json").is_ok());
    }

    #[test]
    fn test_validate_jwks_uri_rejects_http_non_localhost() {
        assert!(matches!(
            validate_jwks_uri("http://auth.example.com/.well-known/jwks.json"),
            Err(JwksError::DisallowedUriScheme(_))
        ));
    }

    #[test]
    fn test_validate_jwks_uri_rejects_non_http_scheme() {
        assert!(matches!(
            validate_jwks_uri("ftp://auth.example.com/.well-known/jwks.json"),
            Err(JwksError::DisallowedUriScheme(_))
        ));
    }

    #[test]
    fn test_validate_jwks_uri_rejects_unparseable_uri() {
        assert!(matches!(
            validate_jwks_uri("not a uri"),
            Err(JwksError::DisallowedUriScheme(_))
        ));
    }
}

#[cfg(test)]
mod test_mock {
    use std::sync::Arc;

    use super::mock::{MockJwksProvider, MockJwksResponse};
    use super::{JwksError, JwksProvider};

    fn dummy_client() -> reqwest::Client {
        reqwest::Client::new()
    }

    #[tokio::test]
    async fn test_mock_returns_responses_in_order() {
        let provider = MockJwksProvider::new(vec![
            Arc::new(|| Ok(vec![("key-1".to_string(), dummy_decoding_key())])) as MockJwksResponse,
            Arc::new(|| Ok(vec![("key-2".to_string(), dummy_decoding_key())])) as MockJwksResponse,
        ]);

        let first = provider
            .fetch_jwks(dummy_client(), "http://ignored".to_string())
            .await
            .unwrap();
        let second = provider
            .fetch_jwks(dummy_client(), "http://ignored".to_string())
            .await
            .unwrap();

        assert_eq!(first[0].0, "key-1");
        assert_eq!(second[0].0, "key-2");
    }

    #[tokio::test]
    async fn test_mock_reuses_last_response_when_exhausted() {
        let provider = MockJwksProvider::new(vec![Arc::new(|| {
            Ok(vec![("key-1".to_string(), dummy_decoding_key())])
        }) as MockJwksResponse]);

        for _ in 0..3 {
            let keys = provider
                .fetch_jwks(dummy_client(), "http://ignored".to_string())
                .await
                .unwrap();
            assert_eq!(keys[0].0, "key-1");
        }
    }

    #[tokio::test]
    async fn test_mock_can_return_errors() {
        let provider = MockJwksProvider::new(vec![Arc::new(|| {
            Err(JwksError::ErrorStatus(
                reqwest::StatusCode::INTERNAL_SERVER_ERROR,
            ))
        }) as MockJwksResponse]);

        assert!(matches!(
            provider
                .fetch_jwks(dummy_client(), "http://ignored".to_string())
                .await
                .unwrap_err(),
            JwksError::ErrorStatus(_)
        ));
    }

    #[test]
    #[should_panic(expected = "MockJwksProvider requires at least one response")]
    fn test_mock_panics_on_empty_responses() {
        MockJwksProvider::new(vec![]);
    }

    /// Returns a valid `DecodingKey` for use in mock responses where key material
    /// is not exercised (only the `kid` string is checked).
    fn dummy_decoding_key() -> jsonwebtoken::DecodingKey {
        use jsonwebtoken::DecodingKey;

        const RSA_N: &str = "rguts8kRMaf55JSXIfp0Nt1GYOEZeCaAwGiegSPZuKOh7OqzjhGCBayaNmoPEB86gRGKzpW9cdsSHAholAsQXiQYrMNZMgcDLT3TE_4n5KeL3zEsFL-lsV9QndM2chpZ3_U_jzreOMz_U0JbBS6WsNU2HYwsFW-vU69fTVqChKf4SfbM9rfV1I7zLZxEpYKyVIoSMXydTE90TuoPtXMtOLgRuBjp5En_szHqY3nkF-VbB9Fq1mhu1LsnuynbxTWErwWf43rqGt9BFyvouwyzdfOD9H69hxw5i09De9dXLAU7nl5C9AsOs_67DY4VFznMWRQ6AoFEhd-j5NzopVa3nQ";
        const RSA_E: &str = "AQAB";
        DecodingKey::from_rsa_components(RSA_N, RSA_E).unwrap()
    }
}