use std::sync::Arc;

use axum::{RequestPartsExt, extract::FromRequestParts};
use axum_extra::{
    TypedHeader,
    headers::{Authorization, authorization::Bearer},
};
use http::{StatusCode, request::Parts};
use jsonwebtoken::{DecodingKey, jwk::JwkSet};
use serde::Deserialize;
use services::User;
use tokio::sync::RwLock;
use tracing::{info, warn};

use crate::AppState;

#[derive(Clone, Debug, Deserialize)]
struct Claims {
    id: String,
}

impl FromRequestParts<AppState> for User {
    type Rejection = (StatusCode, &'static str);

    async fn from_request_parts(
        parts: &mut Parts,
        state: &AppState,
    ) -> Result<Self, Self::Rejection> {
        if let Some(user) = &state.test_user {
            return Ok(user.clone());
        }

        let TypedHeader(Authorization(bearer)) = parts
            .extract::<TypedHeader<Authorization<Bearer>>>()
            .await
            .map_err(|_| (StatusCode::UNAUTHORIZED, "Authorization header missing"))?;
        let raw_jwt = bearer.token();

        let header = jsonwebtoken::decode_header(raw_jwt)
            .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid token"))?;

        let jwk = state
            .jwks_cache
            .get_key(&state.constants.auth_api_base, header.kid.as_deref())
            .await
            .map_err(|e| {
                warn!("Failed to get key: {:?}", e);
                (StatusCode::INTERNAL_SERVER_ERROR, "Failed to get key")
            })?
            .ok_or((StatusCode::UNAUTHORIZED, "Invalid token"))?;

        let key = DecodingKey::from_jwk(&jwk).map_err(|_| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                "Failed to create decoding key",
            )
        })?;

        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::EdDSA);
        validation.set_audience(&[&state.constants.self_api_base]);

        let claims: Claims = jsonwebtoken::decode(raw_jwt, &key, &validation)
            .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid token"))?
            .claims;

        Ok(User { id: claims.id })
    }
}

#[derive(Clone)]
pub struct JwksCache {
    inner: Arc<RwLock<Option<JwkSet>>>,
}

impl Default for JwksCache {
    fn default() -> Self {
        Self {
            inner: Arc::new(RwLock::new(None)),
        }
    }
}

impl JwksCache {
    async fn get_key(
        &self,
        auth_api_base: &str,
        kid: Option<&str>,
    ) -> Result<Option<jsonwebtoken::jwk::Jwk>, reqwest::Error> {
        // Fast path: cached lookup
        let guard = self.inner.read().await;
        if let Some(jwks) = &*guard
            && let Some(jwk) = self.find_jwk(jwks, kid)
        {
            return Ok(Some(jwk.clone()));
        }
        drop(guard);

        // Slow path: single refresh
        let mut guard = self.inner.write().await;

        // Double-check in case another task refreshed already
        if let Some(jwks) = &*guard
            && let Some(jwk) = self.find_jwk(jwks, kid)
        {
            return Ok(Some(jwk.clone()));
        }

        info!("Refreshing JWKS from {}", auth_api_base);
        let fresh = reqwest::get(format!("{}/api/auth/jwks", auth_api_base))
            .await?
            .json::<JwkSet>()
            .await?;

        let found = self.find_jwk(&fresh, kid);
        *guard = Some(fresh);

        Ok(found)
    }

    fn find_jwk(&self, jwks: &JwkSet, kid: Option<&str>) -> Option<jsonwebtoken::jwk::Jwk> {
        jwks.keys
            .iter()
            .find(|j| j.common.key_id.as_deref() == kid)
            .cloned()
    }
}