//! @docs ARCHITECTURE:Security
//!
//! ### AI Assist Note
//! **Auth Middleware**: Orchestrates the validation of **Bearer Tokens**
//! for the Tadpole OS engine. Implements **Constant-Time Comparison**
//! (`constant_time_eq`) to mask timing-based side-channel attacks on
//! the `NEURAL_TOKEN`. Supports **Dual-Mechanism Validation**: standard
//! `Authorization` headers for REST, and `Sec-WebSocket-Protocol`
//! (bearer.<token>) for browser-based WebSocket upgrades (AUTH-01).
//!
//! ### 🔍 Debugging & Observability
//! - **Failure Path**: Missing `Authorization` header, invalid token
//!   checksum, or malformed WebSocket subprotocol strings. Triggers
//!   `401 Unauthorized`.
//! - **Trace Scope**: `server-rs::middleware::auth`

use crate::AppState;
use axum::{
    body::Body,
    extract::State,
    http::{Request, StatusCode},
    middleware::Next,
    response::Response,
};
use std::sync::Arc;

use subtle::ConstantTimeEq;

/// Constant-time string comparison to prevent timing-based side-channel attacks.
///
/// ### 🔒 Security: Constant-Time Comparison (AUTH-01)
/// Standard string equality checks return `false` as soon as they find the
/// first differing byte. An attacker can use this timing information to
/// guess a token one character at a time.
///
/// This implementation uses the `subtle` crate to ensure that the execution
/// time is deterministic relative to the input length, preventing
/// optimizer-induced early returns.
pub(crate) fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    a.ct_eq(b).into()
}
fn is_valid_access_token(token: &str, state: &AppState) -> bool {
    if token.is_empty() {
        return false;
    }
    (!state.security.deploy_token.is_empty()
        && constant_time_eq(token.as_bytes(), state.security.deploy_token.as_bytes()))
        || (!state.security.admin_token.is_empty()
            && constant_time_eq(token.as_bytes(), state.security.admin_token.as_bytes()))
}

/// Marker struct to indicate that the request has been authenticated via header or subprotocol before upgrade.
#[derive(Clone, Copy, Debug)]
pub struct PreAuthenticated;

/// Middleware to validate the Bearer token.
/// Supports two mechanisms:
/// 1. Standard `Authorization: Bearer <token>` header (REST endpoints)
/// 2. `Sec-WebSocket-Protocol: bearer.<token>` header (browser WebSocket upgrades)
pub async fn validate_token(
    State(state): State<Arc<AppState>>,
    mut req: Request<Body>,
    next: Next,
) -> Result<Response, StatusCode> {
    // Check for standard Authorization header first
    let auth_header = req
        .headers()
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|val| val.to_str().ok());

    if let Some(auth_str) = auth_header {
        if let Some(token) = auth_str.strip_prefix("Bearer ") {
            // Guard against empty token bypass (e.g. if deploy_token is empty/unset or input token is empty)
            if is_valid_access_token(token, &state) {
                req.extensions_mut().insert(PreAuthenticated);
                return Ok(next.run(req).await);
            } else {
                tracing::warn!("🚫 Invalid token provided in Authorization header");
                return Err(StatusCode::UNAUTHORIZED);
            }
        }
    }

    // Fallback: check Sec-WebSocket-Protocol for browser WS connections
    // Browsers cannot set Authorization headers on WebSocket upgrade requests,
    // so the frontend sends the token as a subprotocol: "bearer.<token>"
    let is_ws_upgrade = req
        .headers()
        .get(axum::http::header::UPGRADE)
        .and_then(|v| v.to_str().ok())
        .map(|v| v.eq_ignore_ascii_case("websocket"))
        .unwrap_or(false);

    if is_ws_upgrade {
        let proto_header = req
            .headers()
            .get("sec-websocket-protocol")
            .and_then(|v| v.to_str().ok())
            .unwrap_or("");

        let mut has_bearer = false;
        let mut is_bearer_valid = false;

        // SEC-01: Split the comma-separated list of protocols
        // Browsers often combine multiple subprotocols in one header
        for protocol in proto_header.split(',') {
            let protocol = protocol.trim();
            if let Some(token) = protocol.strip_prefix("bearer.") {
                has_bearer = true;
                // Guard against empty token bypass
                if is_valid_access_token(token, &state) {
                    is_bearer_valid = true;
                }
            }
        }

        if has_bearer {
            if is_bearer_valid {
                req.extensions_mut().insert(PreAuthenticated);
                return Ok(next.run(req).await);
            } else {
                tracing::warn!(
                    "🚫 Unauthorized WebSocket upgrade: invalid bearer subprotocol token"
                );
                return Err(StatusCode::UNAUTHORIZED);
            }
        } else {
            // Upgrade request doesn't request bearer.<token> (e.g. it requests 'tadpole-pulse-v1').
            // We allow the upgrade to proceed, but since PreAuthenticated is NOT in extensions,
            // the WebSocket handler will require post-connect authentication.
            return Ok(next.run(req).await);
        }
    } else {
        tracing::warn!("🚫 Missing or malformed Authorization header");
    }

    Err(StatusCode::UNAUTHORIZED)
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{
        body::Body,
        http::{header, Request, StatusCode},
        middleware::from_fn_with_state,
        routing::get,
        Router,
    };
    use tower::ServiceExt;

    async fn dummy_handler() -> StatusCode {
        StatusCode::OK
    }

    #[tokio::test]
    async fn test_auth_bearer_success() {
        let state = Arc::new(AppState::new_minimal_mock().await);
        let app = Router::new()
            .route("/", get(dummy_handler))
            .layer(from_fn_with_state(state, validate_token));

        let req = Request::builder()
            .uri("/")
            .header(header::AUTHORIZATION, "Bearer test-token")
            .body(Body::empty())
            .unwrap();

        let res = app.oneshot(req).await.unwrap();
        assert_eq!(res.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_auth_bearer_invalid() {
        let state = Arc::new(AppState::new_minimal_mock().await);
        let app = Router::new()
            .route("/", get(dummy_handler))
            .layer(from_fn_with_state(state, validate_token));

        let req = Request::builder()
            .uri("/")
            .header(header::AUTHORIZATION, "Bearer wrong-token")
            .body(Body::empty())
            .unwrap();

        let res = app.oneshot(req).await.unwrap();
        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn test_auth_websocket_protocol_success() {
        let state = Arc::new(AppState::new_minimal_mock().await);
        let app = Router::new()
            .route("/", get(dummy_handler))
            .layer(from_fn_with_state(state, validate_token));

        let req = Request::builder()
            .uri("/")
            .header(header::UPGRADE, "websocket")
            .header("sec-websocket-protocol", "bearer.test-token")
            .body(Body::empty())
            .unwrap();

        let res = app.oneshot(req).await.unwrap();
        assert_eq!(res.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_auth_websocket_protocol_invalid_bearer() {
        let state = Arc::new(AppState::new_minimal_mock().await);
        let app = Router::new()
            .route("/", get(dummy_handler))
            .layer(from_fn_with_state(state, validate_token));

        // Invalid bearer token
        let req = Request::builder()
            .uri("/")
            .header(header::UPGRADE, "websocket")
            .header("sec-websocket-protocol", "bearer.wrong-token")
            .body(Body::empty())
            .unwrap();

        let res = app.oneshot(req).await.unwrap();
        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn test_auth_websocket_protocol_post_connect() {
        let state = Arc::new(AppState::new_minimal_mock().await);
        let app = Router::new()
            .route("/", get(dummy_handler))
            .layer(from_fn_with_state(state, validate_token));

        // No bearer token in protocol (post-connect auth mode)
        let req = Request::builder()
            .uri("/")
            .header(header::UPGRADE, "websocket")
            .header("sec-websocket-protocol", "tadpole-pulse-v1")
            .body(Body::empty())
            .unwrap();

        let res = app.oneshot(req).await.unwrap();
        assert_eq!(res.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_auth_missing_header() {
        let state = Arc::new(AppState::new_minimal_mock().await);
        let app = Router::new()
            .route("/", get(dummy_handler))
            .layer(from_fn_with_state(state, validate_token));

        let req = Request::builder().uri("/").body(Body::empty()).unwrap();

        let res = app.oneshot(req).await.unwrap();
        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn test_auth_empty_token() {
        let state = Arc::new(AppState::new_minimal_mock().await);
        let app = Router::new()
            .route("/", get(dummy_handler))
            .layer(from_fn_with_state(state, validate_token));

        let req = Request::builder()
            .uri("/")
            .header(header::AUTHORIZATION, "Bearer ")
            .body(Body::empty())
            .unwrap();

        let res = app.oneshot(req).await.unwrap();
        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
    }
}

/// Extractor that requires administrative credentials (admin_token).
#[derive(Clone, Copy, Debug)]
pub struct RequireAdmin;

impl<S> axum::extract::FromRequestParts<S> for RequireAdmin
where
    Arc<AppState>: axum::extract::FromRef<S>,
    S: Send + Sync,
{
    type Rejection = crate::error::AppError;

    async fn from_request_parts(
        parts: &mut axum::http::request::Parts,
        state: &S,
    ) -> Result<Self, Self::Rejection> {
        use axum::extract::FromRef;
        let app_state = Arc::<AppState>::from_ref(state);

        let auth_header = parts
            .headers
            .get(axum::http::header::AUTHORIZATION)
            .and_then(|val| val.to_str().ok());

        if let Some(auth_str) = auth_header {
            if let Some(token) = auth_str.strip_prefix("Bearer ") {
                if !token.is_empty()
                    && !app_state.security.admin_token.is_empty()
                    && constant_time_eq(token.as_bytes(), app_state.security.admin_token.as_bytes())
                {
                    return Ok(RequireAdmin);
                }
            }
        }
        Err(crate::error::AppError::Unauthorized(
            "Administrative privileges required".to_string(),
        ))
    }
}

// Metadata: [auth]
