use std::sync::Arc;
use std::time::{Duration, Instant};

use config::oauth::OAuthConfig;
use jsonwebtoken::jwk::JwkSet;
use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
use tokio::sync::{Mutex, RwLock};

use crate::claims::{AccessClaims, IdClaims, ServiceClaims};
use crate::error::SessionError;

const ACCESS_TYP: &str = "at+jwt";
const SERVICE_TYP: &str = "at+jwt";

struct CachedJwks {
    set: JwkSet,
    fetched_at: Instant,
}

pub struct JwtVerifier {
    issuer: String,
    audience: String,
    chat_audience: String,
    client_id: String,
    session_cookie: String,
    leeway: u64,
    jwks_url: String,
    cache_ttl: Duration,
    min_refresh: Duration,
    http: reqwest::Client,
    cache: RwLock<Option<CachedJwks>>,
    refresh: Mutex<()>,
}

impl JwtVerifier {
    pub fn new(config: &OAuthConfig) -> Result<Arc<Self>, SessionError> {
        if config.issuer.is_empty() {
            return Err(SessionError::Verifier("oauth.issuer is empty".into()));
        }
        let http = reqwest::Client::builder()
            .timeout(Duration::from_millis(config.http_timeout_ms.max(200)))
            .build()
            .map_err(|err| SessionError::Verifier(err.to_string()))?;

        Ok(Arc::new(Self {
            issuer: config.issuer.trim_end_matches('/').to_string(),
            audience: config.audience.clone(),
            chat_audience: config.chat_audience.clone(),
            client_id: config.client_id.clone(),
            session_cookie: config.session_cookie.clone(),
            leeway: config.leeway_secs,
            jwks_url: config.jwks_url(),
            cache_ttl: Duration::from_secs(config.jwks_cache_secs.max(1)),
            min_refresh: Duration::from_secs(config.jwks_min_refresh_secs.max(1)),
            http,
            cache: RwLock::new(None),
            refresh: Mutex::new(()),
        }))
    }

    pub fn jwks_url(&self) -> &str {
        &self.jwks_url
    }

    pub fn session_cookie(&self) -> &str {
        &self.session_cookie
    }

    pub async fn warm_up(&self) {
        if let Err(error) = self.fetch_jwks().await {
            tracing::error!(target: "gitdata::jwks", error = %error, url = %self.jwks_url, "failed to load JWKS");
        }
    }

    pub async fn verify_access(&self, token: &str) -> Result<AccessClaims, SessionError> {
        let (claims, typ) = self.verify::<AccessClaims>(token, &self.audience).await?;
        if typ != ACCESS_TYP {
            return Err(SessionError::Verifier(format!("unexpected typ {typ:?}")));
        }
        if claims.token_type != "access" {
            return Err(SessionError::Verifier(format!(
                "unexpected token_type {:?}",
                claims.token_type
            )));
        }
        if claims.client_id != self.client_id {
            return Err(SessionError::Verifier(format!(
                "unexpected client_id {:?}",
                claims.client_id
            )));
        }
        Ok(claims)
    }

    pub async fn verify_id(&self, token: &str) -> Result<IdClaims, SessionError> {
        let (claims, _) = self.verify::<IdClaims>(token, &self.client_id).await?;
        Ok(claims)
    }

    pub async fn verify_service(
        &self,
        token: &str,
        audience: &str,
    ) -> Result<ServiceClaims, SessionError> {
        let (claims, typ) = self.verify::<ServiceClaims>(token, audience).await?;
        if typ != SERVICE_TYP && typ != "JWT" {
            return Err(SessionError::Verifier(format!("unexpected typ {typ:?}")));
        }
        if claims.token_type != "service" {
            return Err(SessionError::Verifier(format!(
                "unexpected token_type {:?}",
                claims.token_type
            )));
        }
        Ok(claims)
    }

    pub fn chat_audience(&self) -> &str {
        &self.chat_audience
    }

    async fn verify<T: serde::de::DeserializeOwned>(
        &self,
        token: &str,
        audience: &str,
    ) -> Result<(T, String), SessionError> {
        let header =
            decode_header(token).map_err(|err| SessionError::Verifier(err.to_string()))?;
        if header.alg != Algorithm::EdDSA {
            return Err(SessionError::Verifier(format!(
                "unexpected algorithm {:?}",
                header.alg
            )));
        }
        let kid = header
            .kid
            .clone()
            .filter(|value| !value.is_empty())
            .ok_or_else(|| SessionError::Verifier("token has no kid".into()))?;
        let key = self.decoding_key(&kid).await?;

        let mut validation = Validation::new(Algorithm::EdDSA);
        validation.leeway = self.leeway;
        validation.set_issuer(&[self.issuer.as_str()]);
        validation.set_audience(&[audience]);
        validation.required_spec_claims = ["exp", "iss", "aud", "sub"]
            .into_iter()
            .map(str::to_string)
            .collect();

        let data = decode::<T>(token, &key, &validation).map_err(|err| {
            if matches!(
                err.kind(),
                jsonwebtoken::errors::ErrorKind::ExpiredSignature
            ) {
                SessionError::Expired
            } else {
                SessionError::Verifier(err.to_string())
            }
        })?;

        let typ = header.typ.unwrap_or_default();
        Ok((data.claims, typ))
    }

    async fn decoding_key(&self, kid: &str) -> Result<DecodingKey, SessionError> {
        {
            let cache = self.cache.read().await;
            if let Some(cached) = cache.as_ref()
                && cached.fetched_at.elapsed() < self.cache_ttl
                && let Some(key) = decoding_key_for(&cached.set, kid)?
            {
                return Ok(key);
            }
        }

        let _guard = self.refresh.lock().await;
        {
            let cache = self.cache.read().await;
            if let Some(cached) = cache.as_ref()
                && cached.fetched_at.elapsed() < self.min_refresh
            {
                return decoding_key_for(&cached.set, kid)?.ok_or_else(|| {
                    SessionError::Verifier(format!("unknown kid {kid:?} (refresh rate limited)"))
                });
            }
        }

        match self.fetch_jwks().await {
            Ok(set) => {
                let key = decoding_key_for(&set, kid)?;
                *self.cache.write().await = Some(CachedJwks {
                    set,
                    fetched_at: Instant::now(),
                });
                key.ok_or_else(|| SessionError::Verifier(format!("unknown kid {kid:?}")))
            }
            Err(error) => {
                let cache = self.cache.read().await;
                if let Some(cached) = cache.as_ref()
                    && let Some(key) = decoding_key_for(&cached.set, kid)?
                {
                    tracing::error!(target: "gitdata::jwks", error = %error, "JWKS refresh failed; using cached keys");
                    return Ok(key);
                }
                Err(error)
            }
        }
    }

    async fn fetch_jwks(&self) -> Result<JwkSet, SessionError> {
        let response = self
            .http
            .get(&self.jwks_url)
            .send()
            .await
            .map_err(|err| SessionError::Verifier(format!("jwks request failed: {err}")))?;
        if !response.status().is_success() {
            return Err(SessionError::Verifier(format!(
                "jwks endpoint returned {}",
                response.status()
            )));
        }
        response
            .json::<JwkSet>()
            .await
            .map_err(|err| SessionError::Verifier(format!("invalid jwks document: {err}")))
    }
}

fn decoding_key_for(set: &JwkSet, kid: &str) -> Result<Option<DecodingKey>, SessionError> {
    match set.find(kid) {
        Some(jwk) => DecodingKey::from_jwk(jwk)
            .map(Some)
            .map_err(|err| SessionError::Verifier(err.to_string())),
        None => Ok(None),
    }
}