use axum::{extract::Request, http::StatusCode, middleware::Next, response::Response};
use std::sync::{Arc, LazyLock};
use tokio::sync::Semaphore;
use tower_http::request_id::{MakeRequestId, RequestId};
use uuid::Uuid;

/// Maximum concurrent in-flight requests. Responds 429 when at capacity.
const MAX_CONCURRENT: usize = 256;

static SEMAPHORE: LazyLock<Arc<Semaphore>> =
    LazyLock::new(|| Arc::new(Semaphore::new(MAX_CONCURRENT)));

/// Axum middleware that limits concurrent in-flight requests.
/// Returns 429 Too Many Requests when the limit is reached.
pub async fn limit_concurrency(req: Request, next: Next) -> Result<Response, StatusCode> {
    let Ok(_permit) = SEMAPHORE.try_acquire() else {
        return Err(StatusCode::TOO_MANY_REQUESTS);
    };
    Ok(next.run(req).await)
}

/// Generates a UUID v4 for each request and attaches it as `x-request-id`.
#[derive(Clone, Copy)]
pub struct MakeRequestUuid;

impl MakeRequestId for MakeRequestUuid {
    fn make_request_id<B>(&mut self, _: &http::Request<B>) -> Option<RequestId> {
        let id = Uuid::new_v4().to_string();
        let header_val = id.parse().ok()?;
        Some(RequestId::new(header_val))
    }
}