#![cfg(feature = "axum")]

use std::sync::Arc;

use async_trait::async_trait;
use axum::{
    body::Body,
    http::{header, Request, StatusCode},
    routing::get,
    Router,
};
use tower::ServiceExt;

use rok_auth::{
    axum::{AuthLayer, Guard, GuardLayer, JwtGuard, JwtUserProvider},
    Auth, AuthConfig, AuthError, Claims,
};

// ── Helpers ───────────────────────────────────────────────────────────────────

fn test_auth() -> Auth {
    Auth::new(AuthConfig {
        secret: "test-secret-key".to_string(),
        ..Default::default()
    })
}

/// Minimal custom guard: checks `X-Api-Key` header.
struct ApiKeyGuard {
    expected: String,
    claims: Claims,
}

#[async_trait]
impl Guard for ApiKeyGuard {
    async fn authenticate(&self, parts: &axum::http::request::Parts) -> Result<Claims, AuthError> {
        let key = parts.headers.get("x-api-key").and_then(|v| v.to_str().ok());
        if key == Some(self.expected.as_str()) {
            Ok(self.claims.clone())
        } else {
            Err(AuthError::InvalidToken)
        }
    }
}

async fn sub_handler(claims: Claims) -> String {
    claims.sub
}

// ── GuardLayer — custom guard ─────────────────────────────────────────────────

#[tokio::test]
async fn custom_guard_injects_claims_on_valid_key() {
    let guard = ApiKeyGuard {
        expected: "my-secret-key".into(),
        claims: Claims::new("svc-account", vec!["api"]),
    };
    let app = Router::new()
        .route("/", get(sub_handler))
        .layer(AuthLayer::with_guard(guard));

    let resp = app
        .oneshot(
            Request::builder()
                .uri("/")
                .header("x-api-key", "my-secret-key")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(resp.status(), StatusCode::OK);
    let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
        .await
        .unwrap();
    assert_eq!(&body[..], b"svc-account");
}

#[tokio::test]
async fn custom_guard_returns_401_on_invalid_key() {
    let guard = ApiKeyGuard {
        expected: "correct-key".into(),
        claims: Claims::new("svc", Vec::<String>::new()),
    };
    let app = Router::new()
        .route("/", get(sub_handler))
        .layer(AuthLayer::with_guard(guard));

    let resp = app
        .oneshot(
            Request::builder()
                .uri("/")
                .header("x-api-key", "wrong-key")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}

#[tokio::test]
async fn custom_guard_returns_401_when_header_missing() {
    let guard = ApiKeyGuard {
        expected: "key".into(),
        claims: Claims::new("svc", Vec::<String>::new()),
    };
    let app = Router::new()
        .route("/", get(sub_handler))
        .layer(AuthLayer::with_guard(guard));

    let resp = app
        .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}

// ── GuardLayer::optional ──────────────────────────────────────────────────────

// Note: Option<Claims> FromRequestParts uses OptionalClaims internally
// We test that the optional guard passes unauthenticated requests through
#[tokio::test]
async fn optional_guard_passes_unauthenticated_requests_through() {
    let guard = ApiKeyGuard {
        expected: "key".into(),
        claims: Claims::new("svc", Vec::<String>::new()),
    };
    let app = Router::new()
        .route("/", get(|| async { "ok" }))
        .layer(GuardLayer::optional(guard));

    let resp = app
        .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    // Should NOT be 401 — request passes through without auth
    assert_eq!(resp.status(), StatusCode::OK);
}

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

/// Provider that allows any subject.
struct AllowAllProvider;

#[async_trait]
impl JwtUserProvider for AllowAllProvider {
    async fn find_by_subject(&self, sub: &str) -> Option<Claims> {
        Some(Claims::new(sub, vec!["user"]))
    }
}

/// Provider that always denies (user not found / banned).
struct DenyAllProvider;

#[async_trait]
impl JwtUserProvider for DenyAllProvider {
    async fn find_by_subject(&self, _sub: &str) -> Option<Claims> {
        None
    }
}

#[tokio::test]
async fn jwt_guard_accepts_valid_token_with_allowing_provider() {
    let auth = test_auth();
    let token = auth
        .sign(&Claims::new("alice", Vec::<String>::new()))
        .unwrap();
    let guard = JwtGuard::new(Arc::new(auth), AllowAllProvider);

    let app = Router::new()
        .route("/", get(sub_handler))
        .layer(AuthLayer::with_guard(guard));

    let resp = app
        .oneshot(
            Request::builder()
                .uri("/")
                .header(header::AUTHORIZATION, format!("Bearer {}", token))
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(resp.status(), StatusCode::OK);
    let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
        .await
        .unwrap();
    assert_eq!(&body[..], b"alice");
}

#[tokio::test]
async fn jwt_guard_rejects_invalid_jwt() {
    let auth = test_auth();
    let guard = JwtGuard::new(Arc::new(auth), AllowAllProvider);

    let app = Router::new()
        .route("/", get(sub_handler))
        .layer(AuthLayer::with_guard(guard));

    let resp = app
        .oneshot(
            Request::builder()
                .uri("/")
                .header(header::AUTHORIZATION, "Bearer not.a.valid.jwt")
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}

#[tokio::test]
async fn jwt_guard_rejects_when_provider_denies() {
    let auth = test_auth();
    let token = auth
        .sign(&Claims::new("banned-user", Vec::<String>::new()))
        .unwrap();
    let guard = JwtGuard::new(Arc::new(auth), DenyAllProvider);

    let app = Router::new()
        .route("/", get(sub_handler))
        .layer(AuthLayer::with_guard(guard));

    let resp = app
        .oneshot(
            Request::builder()
                .uri("/")
                .header(header::AUTHORIZATION, format!("Bearer {}", token))
                .body(Body::empty())
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}

#[tokio::test]
async fn jwt_guard_rejects_missing_bearer() {
    let auth = test_auth();
    let guard = JwtGuard::new(Arc::new(auth), AllowAllProvider);

    let app = Router::new()
        .route("/", get(sub_handler))
        .layer(AuthLayer::with_guard(guard));

    let resp = app
        .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}

// ── AuthLayer::with_guard factory ────────────────────────────────────────────

#[test]
fn auth_layer_with_guard_builds_guard_layer() {
    let guard = ApiKeyGuard {
        expected: "k".into(),
        claims: Claims::new("x", Vec::<String>::new()),
    };
    let _layer = AuthLayer::with_guard(guard);
    // Just verifies the factory method compiles and constructs without panic
}
