//! Bearer-token authentication for incoming MCP HTTP requests.
//!
//! LLM clients (Claude Desktop / Cursor / ChatGPT / etc.) can't run
//! Keycloak OIDC, so the MCP server uses opaque static tokens minted
//! per-tenant via the agent's `POST /api/v1/mcp-tokens` endpoint.
//!
//! Flow per request:
//! 1. Extract `Authorization: Bearer <token>`. Missing → 401.
//! 2. SHA-256 hash the token.
//! 3. Look up the hash in `<prefix>__admin.mcp_tokens`. Missing or
//!    revoked → 401.
//! 4. Fire-and-forget update of `last_used_at` so the dashboard can
//!    show staleness without blocking the handler.
//! 5. Stash the tenant_id in [`TENANT_ID`] (a `tokio::task_local`) so
//!    the MCP tool handlers can read it without modifying rmcp's
//!    handler signatures.
//!
//! The `task_local` is scoped around the inner service call via
//! [`bearer_auth`], so every handler invoked downstream sees the
//! tenant_id without us having to thread it through the macro-
//! generated tool router.

use axum::body::Body;
use axum::extract::{Request, State};
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use mongodb::bson::doc;
use sha2::{Digest, Sha256};

use crate::database::DatabasePool;

tokio::task_local! {
    /// Tenant id resolved from the bearer for this request. Set by
    /// [`bearer_auth`] before the inner service runs; read by the
    /// MCP tool handlers via [`current_tenant_id`].
    pub static TENANT_ID: String;
}

/// Mongo collection name in `<prefix>__admin`.
const COLLECTION: &str = "mcp_tokens";

/// Returns the tenant_id set by the auth middleware. `None` outside a
/// request scope (e.g. unit tests that bypass the middleware).
pub fn current_tenant_id() -> Option<String> {
    TENANT_ID.try_with(|s| s.clone()).ok()
}

/// Axum middleware: validate bearer → set [`TENANT_ID`] → call inner.
pub async fn bearer_auth(
    State(pool): State<DatabasePool>,
    request: Request,
    next: Next,
) -> Response {
    let Some(token) = extract_bearer(&request) else {
        return (StatusCode::UNAUTHORIZED, "Missing bearer token").into_response();
    };
    if !token.starts_with("mcpt_") {
        return (StatusCode::UNAUTHORIZED, "Invalid token format").into_response();
    }
    let token_hash = sha256_hex(&token);

    let col = pool.admin_db().collection::<TokenLookup>(COLLECTION);
    let found = match col
        .find_one(doc! { "token_hash": &token_hash, "revoked": false })
        .await
    {
        Ok(Some(t)) => t,
        Ok(None) => {
            return (StatusCode::UNAUTHORIZED, "Invalid or revoked token").into_response();
        }
        Err(e) => {
            tracing::error!("MCP token lookup failed: {e}");
            return (StatusCode::INTERNAL_SERVER_ERROR, "Token lookup error").into_response();
        }
    };

    // Fire-and-forget last_used_at update — never block the handler.
    let col2 = pool.admin_db().collection::<TokenLookup>(COLLECTION);
    let hash_for_update = token_hash.clone();
    tokio::spawn(async move {
        let _ = col2
            .update_one(
                doc! { "token_hash": &hash_for_update },
                doc! { "$set": { "last_used_at": mongodb::bson::DateTime::now() } },
            )
            .await;
    });

    let tenant_id = found.tenant_id;
    let inner = next.run(request);
    TENANT_ID.scope(tenant_id, inner).await
}

/// Bare-bones projection — we don't need the whole `McpToken` here,
/// just enough to route and confirm validity.
#[derive(serde::Deserialize)]
struct TokenLookup {
    tenant_id: String,
}

fn extract_bearer(req: &Request<Body>) -> Option<String> {
    req.headers()
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.strip_prefix("Bearer "))
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
}

fn sha256_hex(s: &str) -> String {
    let mut h = Sha256::new();
    h.update(s.as_bytes());
    hex::encode(h.finalize())
}

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

    #[test]
    fn sha256_known_value() {
        // python -c 'import hashlib; print(hashlib.sha256(b"mcpt_known").hexdigest())'
        assert_eq!(
            sha256_hex("mcpt_known"),
            "27cf6cf678a44244106863c1c031be8e57b84c2b3019d742f755f8e7afa75dfd"
        );
    }
}
