use axum::{Router, body::Body, response::IntoResponse, routing::get};
use fred::{
    interfaces::ClientLike,
    types::{Builder, Key},
};
use http::{Request, Response, StatusCode, header::AUTHORIZATION};
use std::{io, time::Duration};
use tokio::net::TcpListener;
use tower_rate_limit_fred::{KeyResolver, RateLimitLayer};

// Rate-limit based on the bearer token, returning 401 if the token is missing.
#[derive(Clone)]
struct BearerTokenResolver;

impl KeyResolver for BearerTokenResolver {
    type ErrorBody = Body;

    fn key<T>(&self, req: &Request<T>) -> Result<Option<Key>, Response<Self::ErrorBody>> {
        req.headers()
            .get(AUTHORIZATION)
            .and_then(|v| v.to_str().ok())
            .and_then(|s| s.strip_prefix("Bearer "))
            .map(Key::from)
            .ok_or_else(|| (StatusCode::UNAUTHORIZED, "Missing Bearer token").into_response())
            .map(Some)
    }
}

#[tokio::main]
async fn main() -> io::Result<()> {
    let listener = TcpListener::bind("localhost:8080").await?;

    let pool = Builder::default_centralized()
        .with_performance_config(|config| {
            // highly advised since Fred hangs forever if Redis is unavailable:
            // https://github.com/aembke/fred.rs/issues/353
            config.default_command_timeout = Duration::from_millis(300);
        })
        .build_pool(12)
        .expect("Failed to create Redis pool");
    pool.init().await.expect("Failed to connect to Redis pool");

    let layer = RateLimitLayer::builder()
        .conn(pool)
        .key_resolver(BearerTokenResolver)
        .capacity(5)
        .emission_interval(Duration::from_secs(2))
        .build();

    let router = Router::new()
        .route("/", get(|| async { StatusCode::NO_CONTENT }))
        .route_layer(layer);

    axum::serve(listener, router)
        .with_graceful_shutdown(async {}) // shut down the server so the test terminates lol
        .await
}
