//! Basic Axum example for Configurable Rate Limiting
//!
//! You can configure your router to choose between:
//!  - strategies
//!    - FixedWindow
//!    - SlidingWindow
//!    - TokenBucket
//!  - RateLimiting Style
//!    - Standard (pure rate limiting)
//!    - Throughput (with timeout)
//!    - Latency (with timeout and load-shedding)
//!
//! By default, you get a router which:
//!  - Ensures that no request takes > 500ms
//!  - Uses TokenBucket with capacity requests to 10/second
//!  - Refills Tokens with 10 new tokens/second.
//!
//! We can run `hey` (or your load tool of choice) to verify
//! that this is working correctly as follows:
//!
//! -n: Total requests (100)
//! -c: 1 worker (to ensure the -q limit is global)
//! -q: 20 requests per second
//!
//! We are overloading our limiter to verify enforcement. We
//! are asking it to perform double the requests we support.
//!
//! ```bash
//! hey -n 100 -c 1 -q 20 http://localhost:3000/
//! ```
//!
//! You should see something like 50% 200 and 50% 503 (unavailable) or 408 (timeout).
//!
//! Notes:
//! - The timeout is set as a failsafe to enforce SLAs in case an
//!   inner service is taking too long to complete.
//! - Request limiting is not an exact science, so you'll see some
//!   small variation around 50%.

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use axum::Router;
use axum::error_handling::HandleErrorLayer;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::routing::get;
use clap::Parser;
use clap::ValueEnum;
use shot_limit::FixedWindow;
use shot_limit::Gcra;
use shot_limit::SlidingWindow;
use shot_limit::Strategy;
use shot_limit::TokenBucket;
use tower::BoxError;
use tower::ServiceBuilder;
use tower_governor::GovernorLayer;
use tower_governor::governor::GovernorConfigBuilder;
use tower_governor::key_extractor::GlobalKeyExtractor;
use tower_shot::RateLimitLayer;
use tower_shot::ServiceBuilderExt;
use tower_shot::ShotError;

#[derive(ValueEnum, Clone, Debug)]
enum LimiterType {
    Standard,
    Throughput,
    Latency,
}

#[derive(ValueEnum, Clone, Debug)]
enum StrategyType {
    Fixed,
    Gcra,
    Sliding,
    Bucket,
    Governor,
}

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// The rate limiting algorithm to use
    #[arg(short, long, value_enum, default_value_t = StrategyType::Bucket)]
    strategy: StrategyType,

    /// The type of tower layer to apply
    #[arg(short, long, value_enum, default_value_t = LimiterType::Throughput)]
    limiter: LimiterType,

    /// Requests per second allowed
    #[arg(short, long, default_value_t = 10)]
    capacity: usize,

    /// Increment tokens per period
    #[arg(short, long, default_value_t = 10)]
    increment: usize,

    /// The period over which capacity is replenished (e.g., "1s", "500ms", "1min")
    #[arg(short, long, value_parser = humantime::parse_duration, default_value = "1s")]
    period: Duration,

    /// Maximum time a request can wait for a permit (e.g., "500ms")
    #[arg(short, long, value_parser = humantime::parse_duration, default_value = "500ms")]
    timeout: Duration,

    /// The address to listen on
    #[arg(short, long, default_value = "127.0.0.1:3000")]
    addr: SocketAddr,
}

#[tokio::main]
async fn main() -> Result<(), BoxError> {
    let args = Args::parse();

    let capacity = args.capacity.try_into()?;
    let increment = args.increment.try_into()?;

    // Ensure the Arc is created as the trait object type immediately
    let strategy: Option<Arc<dyn Strategy + Send + Sync>> = match args.strategy {
        StrategyType::Fixed => Some(Arc::new(FixedWindow::new(capacity, args.period))),
        StrategyType::Gcra => Some(Arc::new(Gcra::new(capacity, args.period))),
        StrategyType::Sliding => Some(Arc::new(SlidingWindow::new(capacity, args.period))),
        StrategyType::Bucket => Some(Arc::new(TokenBucket::new(capacity, increment, args.period))),
        StrategyType::Governor => None,
    };

    println!(
        "📡 Strategy: {:?} | Limiter: {:?} | Capacity: {} per {:?} | Timeout: {:?}",
        args.strategy, args.limiter, args.capacity, args.period, args.timeout
    );

    // Setup OTLP metrics
    let meter_provider = init_metrics("axum_managed").await?;

    global::set_meter_provider(meter_provider);

    let mut app = Router::new().route("/", get(|| async { "Hello, Shot!" }));

    let error_layer = HandleErrorLayer::new(handle_shot_error);

    app = match strategy {
        Some(strategy) => match args.limiter {
            // Note: It doesn't matter which type of limiter you choose for Governor,
            // you always get the "low latency" behaviour. That may because I'm not
            // configuring Governor correctly, but I suspect it's becase Governor
            // assumes that inner services are always ready, so "fast fail" is the
            // only behaviour.
            LimiterType::Standard => {
                let layer = RateLimitLayer::new(strategy);
                app.layer(
                    ServiceBuilder::new()
                        .layer(error_layer)
                        .layer(layer)
                        .map_err(BoxError::from),
                )
            }
            LimiterType::Throughput => app.layer(
                ServiceBuilder::new()
                    .layer(error_layer)
                    .throughput_rate_limit(strategy, args.timeout)
                    .map_err(BoxError::from),
            ),
            LimiterType::Latency => app.layer(
                ServiceBuilder::new()
                    .layer(error_layer)
                    .latency_rate_limit(strategy, args.timeout)
                    .map_err(BoxError::from),
            ),
        },
        None => {
            let config = Arc::new(
                GovernorConfigBuilder::default()
                    .period(args.period)
                    .burst_size(capacity.get() as u32)
                    .key_extractor(GlobalKeyExtractor)
                    .finish()
                    .ok_or(BoxError::from("Could not build governor"))?,
            );

            match args.limiter {
                LimiterType::Standard => app.layer(
                    ServiceBuilder::new()
                        .layer(error_layer)
                        .layer(GovernorLayer::new(config))
                        .map_err(BoxError::from),
                ),
                LimiterType::Throughput => app.layer(
                    ServiceBuilder::new()
                        .layer(error_layer)
                        .timeout(args.timeout)
                        .layer(GovernorLayer::new(config))
                        .map_err(BoxError::from),
                ),
                LimiterType::Latency => app.layer(
                    ServiceBuilder::new()
                        .layer(error_layer)
                        .timeout(args.timeout)
                        .load_shed()
                        .layer(GovernorLayer::new(config))
                        .map_err(BoxError::from),
                ),
            }
        }
    };

    let listener = tokio::net::TcpListener::bind(args.addr).await?;

    println!(
        "📡 Listening on http://{}", 
        listener.local_addr()?.to_string()
    );

    axum::serve(listener, app).await?;

    Ok(())
}

/// The signature must match BoxError -> IntoResponse
async fn handle_shot_error(err: BoxError) -> impl IntoResponse {
    if let Ok(shot_err) = err.downcast::<ShotError>() {
        shot_err.into_response()
    } else {
        (StatusCode::INTERNAL_SERVER_ERROR, "Internal Service Error").into_response()
    }
}

use opentelemetry::KeyValue;
use opentelemetry::global;
use opentelemetry_otlp::ExporterBuildError;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::Resource;
use opentelemetry_sdk::metrics::SdkMeterProvider;

/// Configure OTLP exporter for metrics
async fn init_metrics(service_name: &str) -> Result<SdkMeterProvider, ExporterBuildError> {
    // Create resource with service identification
    let resource = Resource::builder()
        .with_attributes(vec![KeyValue::new(
            "service.name",
            service_name.to_string(),
        )])
        .build();

    // Build meter provider with periodic export
    let exporter = opentelemetry_otlp::MetricExporter::builder()
        .with_tonic()
        .with_endpoint(
            std::env::var("OTLP_ENDPOINT").unwrap_or_else(|_| "http://localhost:4317".to_string()),
        )
        .build()?;
    let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
        .with_periodic_exporter(exporter)
        .with_resource(resource)
        .build();
    Ok(meter_provider)
}
