// {{ display_name }} MCP server — generated by mcpify. Do not hand-edit.
//
// Named `oauth2.rs` — the module path itself shadows the extern `oauth2`
// crate within this file's own scope, so it's referenced below via the
// disambiguating leading `::oauth2::...` rather than a bare `oauth2::...`.

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

use async_trait::async_trait;
use serde::Deserialize;

use super::super::auth_strategy::{AuthConfig, AuthStrategy, Credentials};
use super::super::errors::AuthError;

/// Generates a PKCE `code_verifier`/`code_challenge` pair (RFC 7636), for
/// the setup wizard's authorization-code flow.
pub fn create_pkce_pair() -> (String, String) {
    let (challenge, verifier) = ::oauth2::PkceCodeChallenge::new_random_sha256();
    (verifier.secret().clone(), challenge.as_str().to_string())
}

#[derive(Debug, Deserialize)]
struct TokenResponse {
    access_token: String,
    refresh_token: Option<String>,
    expires_in: Option<u64>,
}

#[derive(Debug, Default)]
pub struct OAuth2Strategy;

impl OAuth2Strategy {
    fn from_token_response(
        data: TokenResponse,
        client_id: Option<&str>,
        client_secret: Option<&str>,
        token_url: Option<&str>,
    ) -> Credentials {
        let mut credentials = Credentials::new();
        credentials.insert(
            "authorization_header".to_string(),
            format!("Bearer {}", data.access_token),
        );
        credentials.insert("access_token".to_string(), data.access_token);
        if let Some(refresh_token) = data.refresh_token {
            credentials.insert("refresh_token".to_string(), refresh_token);
        }
        if let Some(expires_in) = data.expires_in {
            let now_ms = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_millis() as u64;
            credentials.insert(
                "expires_at".to_string(),
                (now_ms + expires_in * 1000).to_string(),
            );
        }
        if let Some(client_id) = client_id {
            credentials.insert("client_id".to_string(), client_id.to_string());
        }
        if let Some(client_secret) = client_secret {
            credentials.insert("client_secret".to_string(), client_secret.to_string());
        }
        if let Some(token_url) = token_url {
            credentials.insert("token_url".to_string(), token_url.to_string());
        }
        credentials
    }

    async fn exchange(token_url: &str, params: &[(&str, &str)]) -> anyhow::Result<TokenResponse> {
        let response = reqwest::Client::new()
            .post(token_url)
            .form(params)
            .send()
            .await?
            .error_for_status()?
            .json::<TokenResponse>()
            .await?;
        Ok(response)
    }
}

#[async_trait]
impl AuthStrategy for OAuth2Strategy {
    async fn authenticate(&self, config: &AuthConfig) -> anyhow::Result<Credentials> {
        let client_id = config
            .get("client_id")
            .ok_or_else(|| AuthError::MissingCredentials("client_id, token_url".to_string()))?;
        let token_url = config
            .get("token_url")
            .ok_or_else(|| AuthError::MissingCredentials("client_id, token_url".to_string()))?;

        let mut params = vec![
            ("grant_type", "authorization_code"),
            ("client_id", client_id.as_str()),
        ];
        if let Some(secret) = config.get("client_secret") {
            params.push(("client_secret", secret));
        }
        if let Some(code) = config.get("authorization_code") {
            params.push(("code", code));
        }
        if let Some(redirect_uri) = config.get("redirect_uri") {
            params.push(("redirect_uri", redirect_uri));
        }
        if let Some(verifier) = config.get("code_verifier") {
            params.push(("code_verifier", verifier));
        }

        let response = Self::exchange(token_url, &params).await?;
        Ok(Self::from_token_response(
            response,
            Some(client_id),
            config.get("client_secret").map(String::as_str),
            Some(token_url),
        ))
    }

    async fn refresh_token(&self, credentials: &Credentials) -> anyhow::Result<Credentials> {
        let refresh_token = credentials
            .get("refresh_token")
            .ok_or(AuthError::TokenExpired)?;
        let token_url = credentials
            .get("token_url")
            .ok_or(AuthError::TokenExpired)?;
        let client_id = credentials
            .get("client_id")
            .map(String::as_str)
            .unwrap_or("");

        let mut params = vec![
            ("grant_type", "refresh_token"),
            ("refresh_token", refresh_token.as_str()),
            ("client_id", client_id),
        ];
        if let Some(secret) = credentials.get("client_secret") {
            params.push(("client_secret", secret));
        }

        let response = Self::exchange(token_url, &params).await?;
        Ok(Self::from_token_response(
            response,
            credentials.get("client_id").map(String::as_str),
            credentials.get("client_secret").map(String::as_str),
            Some(token_url),
        ))
    }

    fn validate_credentials(&self, credentials: &Credentials) -> bool {
        if !credentials.contains_key("access_token") {
            return false;
        }
        let Some(expires_at) = credentials.get("expires_at") else {
            return true;
        };
        let Ok(expires_at) = expires_at.parse::<u64>() else {
            return false;
        };
        let now_ms = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;
        now_ms < expires_at
    }
}

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

    #[test]
    fn pkce_pair_produces_a_non_empty_verifier_and_challenge() {
        let (verifier, challenge) = create_pkce_pair();
        assert!(!verifier.is_empty());
        assert!(!challenge.is_empty());
        assert_ne!(verifier, challenge);
    }

    #[test]
    fn validates_a_token_with_no_expiry_as_always_valid() {
        let strategy = OAuth2Strategy;
        let credentials = Credentials::from([("access_token".to_string(), "abc".to_string())]);
        assert!(strategy.validate_credentials(&credentials));
    }

    #[test]
    fn rejects_an_expired_token() {
        let strategy = OAuth2Strategy;
        let credentials = Credentials::from([
            ("access_token".to_string(), "abc".to_string()),
            ("expires_at".to_string(), "1".to_string()),
        ]);
        assert!(!strategy.validate_credentials(&credentials));
    }

    #[test]
    fn rejects_credentials_with_no_access_token() {
        let strategy = OAuth2Strategy;
        assert!(!strategy.validate_credentials(&Credentials::new()));
    }
}
