use std::sync::Arc;

use apollo_router::plugin::Plugin;
use apollo_router::plugin::PluginInit;
use apollo_router::register_plugin;
use apollo_router::services::supergraph;
use http::HeaderMap;
use http::HeaderValue;
use tower::BoxError;
use tower::ServiceBuilder;
use tower::ServiceExt;
use tracing::Span;

const AUTH_BLOCKED_CONTEXT_KEY: &str = "auth_blocked";

mod apollo_client_info;
mod auth_result;
mod check_auth;
mod configuration;
mod token;

use apollo_client_info::APOLLOGRAPHQL_CLIENT_NAME;
use apollo_client_info::APOLLOGRAPHQL_CLIENT_VERSION;
use auth_result::AuthResult;
use check_auth::check_auth;
use configuration::Configuration;
use configuration::RuleConfiguration;
use token::hmac::check::RealHMACChecker;
use token::jwt::check::JWTChecker;
use token::jwt::check::RealJWTChecker;
use token::jwt::configuration::resolve_jwks_uri;

#[derive(Debug)]
struct AuthEnforcement {
    configuration: Configuration,
    jwt_checker: Arc<RealJWTChecker>,
}

#[tracing::instrument(
    level=tracing::Level::ERROR,
    name="check_auth",
    skip_all,
    fields(
        auth.method,
        client.name,
        client.version,
        error.message,
        hmac.sender,
        hmac.recipient,
        jwt.issuer,
        jwt.subject,
        result,
    )
)]
fn check_auth_wrapper<J: JWTChecker>(
    configuration: &RuleConfiguration,
    headers: &HeaderMap<HeaderValue>,
    jwt_checker: &J,
) -> AuthResult {
    if let Some(client_name_header) = headers.get(APOLLOGRAPHQL_CLIENT_NAME) {
        if let Ok(client_name_str) = client_name_header.to_str() {
            Span::current().record("client.name", client_name_str);
        }
    }
    if let Some(client_version_header) = headers.get(APOLLOGRAPHQL_CLIENT_VERSION) {
        if let Ok(client_version_str) = client_version_header.to_str() {
            Span::current().record("client.version", client_version_str);
        }
    }

    let result = check_auth(configuration, headers, &RealHMACChecker {}, jwt_checker);
    match result {
        AuthResult::Ok => {
            Span::current().record("result", "Ok");
        }
        AuthResult::Warning { message } => {
            Span::current().record("error.message", message);
            Span::current().record("result", "Warning");
            tracing::warn!(message);
        }
        AuthResult::Block { message } => {
            Span::current().record("error.message", message);
            Span::current().record("result", "Block");
            tracing::error!(message);
        }
    }
    result
}

#[async_trait::async_trait]
impl Plugin for AuthEnforcement {
    type Config = Configuration;

    async fn new(init: PluginInit<Self::Config>) -> Result<Self, BoxError> {
        tracing::info!(
            "auth_enforcement plugin started, enforcement enabled: {}",
            init.config.enabled
        );
        let jwks_uri = resolve_jwks_uri(&init.config.rules.jwt)?;
        let jwt_checker = Arc::new(
            RealJWTChecker::new(
                &jwks_uri,
                init.config.rules.jwt.jwks_lifespan_secs,
                None,
                None,
            )
            .await?,
        );
        Ok(AuthEnforcement {
            configuration: init.config,
            jwt_checker,
        })
    }

    fn supergraph_service(&self, service: supergraph::BoxService) -> supergraph::BoxService {
        let configuration = self.configuration.clone();
        let jwt_checker = self.jwt_checker.clone();

        ServiceBuilder::new()
            .map_request(move |request: supergraph::Request| {
                if configuration.enabled {
                    let headers = request.supergraph_request.headers();
                    if let AuthResult::Block { .. } =
                        check_auth_wrapper(&configuration.rules, headers, jwt_checker.as_ref())
                    {
                        if let Err(e) = request.context.insert(AUTH_BLOCKED_CONTEXT_KEY, true) {
                            tracing::error!(
                                "Failed to insert {} context value: {}",
                                AUTH_BLOCKED_CONTEXT_KEY,
                                e
                            );
                        }
                    }
                }
                request
            })
            .map_response(|response: supergraph::Response| {
                let context = response.context.clone();
                if context
                    .get::<_, bool>(AUTH_BLOCKED_CONTEXT_KEY)
                    .ok()
                    .flatten()
                    == Some(true)
                {
                    return supergraph::Response::error_builder()
                        .error(
                            apollo_router::graphql::Error::builder()
                                .message("Unauthorized.")
                                .extension_code("UNAUTHORIZED")
                                .build(),
                        )
                        .status_code(http::StatusCode::UNAUTHORIZED)
                        .context(context)
                        .build()
                        .unwrap_or_else(|e| {
                            tracing::error!("Failed to build unauthorized error response: {}", e);
                            response
                        });
                }
                response
            })
            .service(service)
            .boxed()
    }
}

register_plugin!("pde", "auth_enforcement", AuthEnforcement);

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

    #[test]
    fn test_auth_blocked_context_key_is_defined() {
        // Verify the constant is correctly defined and accessible
        assert_eq!(AUTH_BLOCKED_CONTEXT_KEY, "auth_blocked");
    }

    #[test]
    fn test_auth_blocked_context_key_is_not_empty() {
        // Ensure the key is not accidentally set to empty string
        assert!(!AUTH_BLOCKED_CONTEXT_KEY.is_empty());
    }
}