use axum::{
    extract::{ConnectInfo, Request, State},
    http::StatusCode,
    middleware::Next,
    response::{IntoResponse, Response},
};
use std::{
    collections::HashMap,
    net::SocketAddr,
    sync::Arc,
    time::{Duration, Instant},
};
use tokio::sync::Mutex;

#[derive(Clone)]
pub struct RateLimiter {
    limit: u32,
    entries: Arc<Mutex<HashMap<String, (Instant, u32)>>>,
}
impl RateLimiter {
    pub fn new(limit: u32) -> Self {
        Self {
            limit,
            entries: Arc::new(Mutex::new(HashMap::new())),
        }
    }
    async fn allow(&self, key: String) -> bool {
        let mut entries = self.entries.lock().await;
        let value = entries.entry(key).or_insert((Instant::now(), 0));
        if value.0.elapsed() >= Duration::from_secs(60) {
            *value = (Instant::now(), 0);
        }
        if value.1 >= self.limit {
            false
        } else {
            value.1 += 1;
            true
        }
    }
}
pub async fn rate_limit_middleware(
    State(limiter): State<RateLimiter>,
    req: Request,
    next: Next,
) -> Response {
    let key = req
        .extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|v| v.0.ip().to_string())
        .or_else(|| {
            req.headers()
                .get("x-forwarded-for")
                .and_then(|v| v.to_str().ok())
                .and_then(|v| v.split(',').next())
                .map(str::trim)
                .map(str::to_string)
        })
        .unwrap_or_else(|| "unknown".into());
    if limiter.allow(key).await {
        next.run(req).await
    } else {
        (StatusCode::TOO_MANY_REQUESTS, "rate limit exceeded").into_response()
    }
}
