use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    pub sub: String,
    pub exp: usize,
    pub iat: usize,
    pub iss: String,
    pub aud: Audience,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Audience {
    Single(String),
    Multiple(Vec<String>),
}

#[derive(Debug, Deserialize, Clone)]
struct Jwk {
    #[serde(rename = "kty")]
    _kty: String,
    kid: String,
    n: String,
    e: String,
    #[serde(rename = "alg")]
    _alg: Option<String>,
}

#[derive(Debug, Deserialize)]
struct Jwks {
    keys: Vec<Jwk>,
}

pub struct JwksVerifier {
    issuer: String,
    audience: String,
    jwks_uri: String,
    keys: Arc<RwLock<Vec<Jwk>>>,
    client: Client,
}

impl JwksVerifier {
    pub async fn new(issuer: String, audience: String) -> Result<Self, Box<dyn std::error::Error>> {
        let client = Client::new();
        let discovery_url = format!(
            "{}/.well-known/openid-configuration",
            issuer.trim_end_matches('/')
        );
        let config: serde_json::Value = client.get(&discovery_url).send().await?.json().await?;

        let jwks_uri = config["jwks_uri"]
            .as_str()
            .ok_or("Missing jwks_uri in discovery")?
            .to_string();

        let verifier = Self {
            issuer,
            audience,
            jwks_uri,
            keys: Arc::new(RwLock::new(Vec::new())),
            client,
        };

        verifier.refresh_keys().await?;
        Ok(verifier)
    }

    pub async fn refresh_keys(&self) -> Result<(), Box<dyn std::error::Error>> {
        let jwks: Jwks = self.client.get(&self.jwks_uri).send().await?.json().await?;
        let mut keys = self.keys.write().await;
        *keys = jwks.keys;
        Ok(())
    }

    pub async fn verify(&self, token: &str) -> Result<Claims, Box<dyn std::error::Error>> {
        let header = decode_header(token)?;
        let kid = header.kid.ok_or("Missing kid in token header")?;

        let jwk = {
            let keys = self.keys.read().await;
            keys.iter().find(|k| k.kid == kid).cloned()
        };

        let jwk = match jwk {
            Some(key) => key,
            None => {
                self.refresh_keys().await?;
                let keys = self.keys.read().await;
                keys.iter()
                    .find(|k| k.kid == kid)
                    .cloned()
                    .ok_or("Key not found in JWKS")?
            }
        };

        let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?;

        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_issuer(&[self.issuer.clone()]);
        validation.set_audience(&[self.audience.clone()]);
        validation.validate_aud = true;

        let token_data = decode::<Claims>(token, &decoding_key, &validation)?;
        Ok(token_data.claims)
    }
}

pub struct Authenticator {
    client_id: String,
    client_secret: String,
    token_url: String,
    client: Client,
}

impl Authenticator {
    pub async fn new(
        issuer: String,
        client_id: String,
        client_secret: String,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let client = Client::new();
        let discovery_url = format!(
            "{}/.well-known/openid-configuration",
            issuer.trim_end_matches('/')
        );
        let config: serde_json::Value = client.get(&discovery_url).send().await?.json().await?;

        let token_url = config["token_endpoint"]
            .as_str()
            .ok_or("Missing token_endpoint in discovery")?
            .to_string();

        Ok(Self {
            client_id,
            client_secret,
            token_url,
            client,
        })
    }

    pub async fn exchange_code(
        &self,
        code: String,
        redirect_uri: String,
    ) -> Result<serde_json::Value, Box<dyn std::error::Error>> {
        let params = [
            ("grant_type", "authorization_code"),
            ("code", &code),
            ("redirect_uri", &redirect_uri),
            ("client_id", &self.client_id),
            ("client_secret", &self.client_secret),
        ];

        let res = self
            .client
            .post(&self.token_url)
            .form(&params)
            .send()
            .await?
            .json()
            .await?;

        Ok(res)
    }

    pub async fn refresh_token(
        &self,
        refresh_token: String,
    ) -> Result<serde_json::Value, Box<dyn std::error::Error>> {
        let params = [
            ("grant_type", "refresh_token"),
            ("refresh_token", &refresh_token),
            ("client_id", &self.client_id),
            ("client_secret", &self.client_secret),
        ];

        let res = self
            .client
            .post(&self.token_url)
            .form(&params)
            .send()
            .await?
            .json()
            .await?;

        Ok(res)
    }
}