use std::{
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};

use async_trait::async_trait;
use axum::{
    http::{request::Parts, Request},
    response::{IntoResponse, Response},
};
use tower::{Layer, Service};

use crate::{AuthError, Claims};

// ── Guard trait ───────────────────────────────────────────────────────────────

/// Custom authentication guard — extracts and validates credentials from request
/// parts and returns a `Claims` payload on success.
///
/// Implement this trait to plug any credential scheme (API keys, HMAC, custom
/// session cookies, etc.) into the standard `Ctx` / `Claims` extractor pipeline.
///
/// # Example
///
/// ```rust,ignore
/// use async_trait::async_trait;
/// use axum::http::request::Parts;
/// use rok_auth::{Claims, AuthError, axum::{Guard, AuthLayer}};
///
/// struct ApiKeyGuard { key: String }
///
/// #[async_trait]
/// impl Guard for ApiKeyGuard {
///     async fn authenticate(&self, parts: &Parts) -> Result<Claims, AuthError> {
///         let k = parts.headers.get("x-api-key")
///             .and_then(|v| v.to_str().ok())
///             .ok_or(AuthError::InvalidToken)?;
///         if k == self.key { Ok(Claims::new("service", vec![])) }
///         else             { Err(AuthError::InvalidToken) }
///     }
/// }
///
/// let layer = AuthLayer::with_guard(ApiKeyGuard { key: "secret".into() });
/// ```
#[async_trait]
pub trait Guard: Send + Sync + 'static {
    async fn authenticate(&self, parts: &Parts) -> Result<Claims, AuthError>;

    /// Revoke the session associated with `claims`. Default implementation is a no-op
    /// (appropriate for stateless guards).
    async fn revoke(&self, _claims: &Claims) -> Result<(), AuthError> {
        Ok(())
    }
}

// ── GuardLayer ────────────────────────────────────────────────────────────────

/// Tower [`Layer`] that runs a [`Guard`] on every incoming request.
///
/// On success the resolved [`Claims`] are inserted into request extensions so
/// that `Claims::from_request_parts` and `Ctx::from_request_parts` can find them
/// without re-authenticating.
///
/// On failure the guard returns `401 Unauthorized` (or passes through when
/// constructed via [`GuardLayer::optional`]).
pub struct GuardLayer<G: Guard> {
    guard: Arc<G>,
    optional: bool,
}

impl<G: Guard> Clone for GuardLayer<G> {
    fn clone(&self) -> Self {
        Self {
            guard: Arc::clone(&self.guard),
            optional: self.optional,
        }
    }
}

impl<G: Guard> GuardLayer<G> {
    pub fn new(guard: G) -> Self {
        Self {
            guard: Arc::new(guard),
            optional: false,
        }
    }

    /// Build an optional guard layer — unauthenticated requests are passed
    /// through with no `Claims` in extensions (`Ctx.user` will be `None`).
    pub fn optional(guard: G) -> Self {
        Self {
            guard: Arc::new(guard),
            optional: true,
        }
    }
}

impl<G: Guard, S> Layer<S> for GuardLayer<G> {
    type Service = GuardService<S, G>;

    fn layer(&self, inner: S) -> Self::Service {
        GuardService {
            inner,
            guard: Arc::clone(&self.guard),
            optional: self.optional,
        }
    }
}

// ── GuardService ──────────────────────────────────────────────────────────────

pub struct GuardService<S, G: Guard> {
    inner: S,
    guard: Arc<G>,
    optional: bool,
}

impl<S: Clone, G: Guard> Clone for GuardService<S, G> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            guard: Arc::clone(&self.guard),
            optional: self.optional,
        }
    }
}

impl<S, G, ReqBody> Service<Request<ReqBody>> for GuardService<S, G>
where
    S: Service<Request<ReqBody>, Response = Response> + Clone + Send + 'static,
    S::Future: Send + 'static,
    S::Error: Send + 'static,
    G: Guard,
    ReqBody: Send + 'static,
{
    type Response = Response;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Response, S::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
        let guard = Arc::clone(&self.guard);
        let optional = self.optional;
        let mut inner = self.inner.clone();

        Box::pin(async move {
            let (mut parts, body) = req.into_parts();
            match guard.authenticate(&parts).await {
                Ok(claims) => {
                    parts.extensions.insert(claims);
                    inner.call(Request::from_parts(parts, body)).await
                }
                Err(e) => {
                    if optional {
                        inner.call(Request::from_parts(parts, body)).await
                    } else {
                        Ok(e.into_response())
                    }
                }
            }
        })
    }
}

// ── JwtUserProvider ───────────────────────────────────────────────────────────

/// Pluggable user-lookup hook for [`JwtGuard`].
///
/// Implement this to add database-level validation to JWT authentication —
/// for example, to reject tokens issued to banned or deleted users.
#[async_trait]
pub trait JwtUserProvider: Send + Sync + 'static {
    /// Return `Some(claims)` to allow the request, or `None` to deny it.
    async fn find_by_subject(&self, subject: &str) -> Option<Claims>;
}

// ── JwtGuard ──────────────────────────────────────────────────────────────────

/// A [`Guard`] that validates a Bearer JWT and then delegates to a
/// [`JwtUserProvider`] for additional user-level checks.
pub struct JwtGuard<P: JwtUserProvider> {
    auth: Arc<crate::Auth>,
    provider: Arc<P>,
}

impl<P: JwtUserProvider> JwtGuard<P> {
    pub fn new(auth: Arc<crate::Auth>, provider: P) -> Self {
        Self {
            auth,
            provider: Arc::new(provider),
        }
    }
}

#[async_trait]
impl<P: JwtUserProvider> Guard for JwtGuard<P> {
    async fn authenticate(&self, parts: &Parts) -> Result<Claims, AuthError> {
        let token = bearer_token(parts).ok_or(AuthError::InvalidToken)?;
        let claims = self.auth.verify(&token)?;
        self.provider
            .find_by_subject(&claims.sub)
            .await
            .ok_or(AuthError::InvalidToken)
    }
}

fn bearer_token(parts: &Parts) -> Option<String> {
    parts
        .headers
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer ").map(String::from))
}
