//! Rate limit tester for the RugCheck.xyz summary API.
//!
//! Runs repeated sweep cycles over a configurable duration (default 60 s).
//! Each sweep doubles concurrency until the rate limit is hit or the max is
//! reached. Results are collected across all sweeps and averaged to produce a
//! stable rate-limit estimate.
//!
//! # Measured rate limit (2026-02-19, unauthenticated, single IP)
//!
//! 11 sweeps over 58 s against `GET /v1/tokens/{mint}/report/summary`:
//!
//! | Metric                  | Value      |
//! |-------------------------|------------|
//! | Min measured limit      | 6.1 req/s  |
//! | Max measured limit      | 10.3 req/s |
//! | **Avg measured limit**  | **7.0 req/s** |
//! | **Recommended safe rate** | **≤ 5.6 req/s** (avg × 0.8) |
//!
//! The server appears to use a token bucket of ~3 tokens that refills roughly
//! every 0.5 s (≈ 6 req/s sustained). First-call bursts can reach ~10 req/s
//! when the bucket is full after an idle period. Stay at or below **5 req/s**
//! for reliable zero-429 operation without authentication.
//!
//! # Usage
//! ```bash
//! cargo run --example rugcheck_rate_limit
//! cargo run --example rugcheck_rate_limit -- --requests 10 --max-concurrency 16 --duration 60
//! ```

use anyhow::Result;
use colored::*;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;
use tokio::task::JoinSet;

const RUGCHECK_URL: &str = "https://api.rugcheck.xyz/v1/tokens/{}/report/summary";

/// Distinct well-known Solana token mints used as request targets.
/// Using varied addresses avoids any server-side caching effects.
const TEST_TOKENS: &[&str] = &[
    "So11111111111111111111111111111111111111112", // Wrapped SOL
    "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v", // USDC
    "Es9vMFrzaCERmJfrF4H2FYD4KCoNkY11McCe8BenwNYB", // USDT
    "mSoLzYCxHdYgdzU16g5QSh3i5K3z3KZK7ytfqcJm7So", // mSOL
    "7vfCXTUXx5WJV5JADk17DUJ4ksgau7utNKj4b963voxs", // ETH (Wormhole)
    "DezXAZ8z7PnrnRJjz3wXBoRgixCa6xjnB7YaB1pPB263", // BONK
    "EKpQGSJtjMFqKZ9KQanSqYXRcF8fBopzLHYxdM65zcjm", // WIF
    "JUPyiwrYJFskUPiHa7hkeR8VUtAeFoSYbKedZNsDvCN", // JUP
    "HZ1JovNiVvGrGNiiYvEozEVgZ58xaU3RKwX8eACQBCt3", // PYTH
    "4k3Dyjzvzp8eMZWUXbBCjEvwSkkk59S5iCNLY3QrkX6R", // RAY
];

/// Cooldown between sweeps so the server token-bucket can partially refill.
const INTER_SWEEP_COOLDOWN: Duration = Duration::from_secs(5);
/// Cooldown between concurrency rounds within a single sweep.
const INTER_ROUND_COOLDOWN: Duration = Duration::from_millis(1_000);

// ── Types ─────────────────────────────────────────────────────────────────────

struct RequestResult {
    status: u16,
    latency: Duration,
}

struct RoundStats {
    ok: usize,
    rate_limited: usize,
    total: usize,
    elapsed: Duration,
}

impl RoundStats {
    fn hit_limit(&self) -> bool {
        self.rate_limited * 2 > self.total
    }
    /// Successful requests per second during this round.
    fn ok_rps(&self) -> f64 {
        self.ok as f64 / self.elapsed.as_secs_f64()
    }
}

// ── Fetch ─────────────────────────────────────────────────────────────────────

async fn fetch_summary(client: Arc<reqwest::Client>, token: String) -> RequestResult {
    let url = RUGCHECK_URL.replace("{}", &token);
    let start = Instant::now();
    match client.get(&url).send().await {
        Ok(resp) => RequestResult {
            status: resp.status().as_u16(),
            latency: start.elapsed(),
        },
        Err(_) => RequestResult {
            status: 0,
            latency: start.elapsed(),
        },
    }
}

// ── Burst runner ──────────────────────────────────────────────────────────────

/// Fire `total` requests with at most `concurrency` in-flight simultaneously.
async fn run_burst(
    client: Arc<reqwest::Client>,
    concurrency: usize,
    total: usize,
) -> (Vec<RequestResult>, Duration) {
    let sem = Arc::new(Semaphore::new(concurrency));
    let mut join_set: JoinSet<RequestResult> = JoinSet::new();
    let wall_start = Instant::now();

    for idx in 0..total {
        let permit = sem.clone().acquire_owned().await.unwrap();
        let c = client.clone();
        let token = TEST_TOKENS[idx % TEST_TOKENS.len()].to_string();
        join_set.spawn(async move {
            let result = fetch_summary(c, token).await;
            drop(permit);
            result
        });
    }

    let mut results = Vec::with_capacity(total);
    while let Some(res) = join_set.join_next().await {
        if let Ok(r) = res {
            results.push(r);
        }
    }

    (results, wall_start.elapsed())
}

// ── Table row ─────────────────────────────────────────────────────────────────

fn print_row(concurrency: usize, results: &[RequestResult], elapsed: Duration) -> RoundStats {
    let n = results.len();
    let ok = results
        .iter()
        .filter(|r| (200..300).contains(&r.status))
        .count();
    let rate_limited = results.iter().filter(|r| r.status == 429).count();
    let errors = results
        .iter()
        .filter(|r| r.status == 0 || r.status >= 500)
        .count();

    let mut latencies: Vec<u128> = results.iter().map(|r| r.latency.as_millis()).collect();
    latencies.sort_unstable();
    let avg_ms = if n > 0 {
        latencies.iter().sum::<u128>() / n as u128
    } else {
        0
    };
    let p95_idx = if n > 0 { (n * 95 / 100).min(n - 1) } else { 0 };
    let p95_ms = if n > 0 { latencies[p95_idx] } else { 0 };
    let rps = n as f64 / elapsed.as_secs_f64();

    let ok_col = if ok == n {
        format!("{ok}").green()
    } else if ok > 0 {
        format!("{ok}").yellow()
    } else {
        format!("{ok}").red()
    };
    let rl_col = if rate_limited > 0 {
        format!("{rate_limited}").bright_red().bold()
    } else {
        format!("{rate_limited}").dimmed()
    };
    let err_col = if errors > 0 {
        format!("{errors}").red()
    } else {
        format!("{errors}").dimmed()
    };

    // Extra width on colored columns absorbs embedded ANSI codes.
    print!("{:<12}", concurrency);
    print!("{:<8}", n);
    print!("{:<20}", ok_col);
    print!("{:<22}", rl_col);
    print!("{:<20}", err_col);
    print!("{:<10}", avg_ms);
    print!("{:<10}", p95_ms);
    println!("{:.1}", rps);

    RoundStats {
        ok,
        rate_limited,
        total: n,
        elapsed,
    }
}

fn print_table_header() {
    println!(
        "{:<12}{:<8}{:<20}{:<22}{:<20}{:<10}{:<10}{}",
        "Concurrency", "Sent", "2xx (OK)", "429 (Rate-Ltd)", "Err", "Avg ms", "P95 ms", "Req/s"
    );
    println!("{}", "─".repeat(90).dimmed());
}

// ── Sweep ─────────────────────────────────────────────────────────────────────

/// Run one full concurrency sweep. Returns the collected round stats and the
/// measured ok req/s at the point rate limiting was first observed (or the
/// peak ok req/s if no rate limiting occurred).
async fn run_sweep(
    client: Arc<reqwest::Client>,
    requests_per_round: usize,
    max_concurrency: usize,
) -> (Vec<RoundStats>, f64) {
    let mut rounds: Vec<RoundStats> = Vec::new();
    let mut concurrency = 1usize;

    loop {
        let c = concurrency.min(max_concurrency);
        let total = requests_per_round.max(c);

        let (results, elapsed) = run_burst(client.clone(), c, total).await;
        let stats = print_row(c, &results, elapsed);
        let hit_limit = stats.hit_limit();
        rounds.push(stats);

        if hit_limit {
            println!(
                "  {}",
                format!("↳ rate limit triggered at concurrency={c}").bright_red()
            );
            break;
        }

        if concurrency >= max_concurrency {
            break;
        }

        tokio::time::sleep(INTER_ROUND_COOLDOWN).await;
        concurrency = (concurrency * 2).min(max_concurrency);
    }

    // Measured limit: ok_rps of the rate-limited round, else peak clean ok_rps.
    let measured = rounds
        .iter()
        .find(|r| r.hit_limit())
        .map(|r| r.ok_rps())
        .unwrap_or_else(|| rounds.iter().map(|r| r.ok_rps()).fold(0.0_f64, f64::max));

    (rounds, measured)
}

// ── Entry point ───────────────────────────────────────────────────────────────

#[tokio::main]
async fn main() -> Result<()> {
    let mut requests_per_round: usize = 10;
    let mut max_concurrency: usize = 16;
    let mut duration_secs: u64 = 60;

    let args: Vec<String> = std::env::args().collect();
    let mut i = 1;
    while i < args.len() {
        match args[i].as_str() {
            "--requests" | "-n" if i + 1 < args.len() => {
                requests_per_round = args[i + 1].parse().unwrap_or(requests_per_round);
                i += 2;
            }
            "--max-concurrency" | "-c" if i + 1 < args.len() => {
                max_concurrency = args[i + 1].parse().unwrap_or(max_concurrency);
                i += 2;
            }
            "--duration" | "-d" if i + 1 < args.len() => {
                duration_secs = args[i + 1].parse().unwrap_or(duration_secs);
                i += 2;
            }
            _ => i += 1,
        }
    }

    println!(
        "{}",
        "=== RugCheck.xyz Rate Limit Tester ==="
            .bright_cyan()
            .bold()
    );
    println!("Endpoint : {}", RUGCHECK_URL.dimmed());
    println!(
        "Requests : {} per round",
        requests_per_round.to_string().bright_white()
    );
    println!("Max conc : {}", max_concurrency.to_string().bright_white());
    println!("Duration : {} s", duration_secs.to_string().bright_white());
    println!(
        "Strategy : double concurrency each round | {} s inter-sweep cooldown",
        INTER_SWEEP_COOLDOWN.as_secs()
    );
    println!();

    let client = Arc::new(
        reqwest::Client::builder()
            .timeout(Duration::from_secs(15))
            .build()?,
    );

    let run_start = Instant::now();
    let total_duration = Duration::from_secs(duration_secs);
    let mut sweep_num = 0usize;
    let mut measured_limits: Vec<f64> = Vec::new();

    while run_start.elapsed() < total_duration {
        sweep_num += 1;
        let elapsed_total = run_start.elapsed();
        let remaining = total_duration.saturating_sub(elapsed_total);

        println!(
            "{}",
            format!(
                "── Sweep #{sweep_num}  (elapsed {:.0}s / {duration_secs}s, {:.0}s remaining) ──",
                elapsed_total.as_secs_f64(),
                remaining.as_secs_f64(),
            )
            .bright_cyan()
        );
        print_table_header();

        let (_, measured) = run_sweep(client.clone(), requests_per_round, max_concurrency).await;
        measured_limits.push(measured);

        println!("  {} {:.1} req/s", "Sweep measured limit:".bold(), measured,);
        println!();

        // Cooldown before next sweep, skip if time is up.
        if run_start.elapsed() + INTER_SWEEP_COOLDOWN < total_duration {
            println!(
                "{}",
                format!(
                    "  Cooling down {}s before next sweep…",
                    INTER_SWEEP_COOLDOWN.as_secs()
                )
                .dimmed()
            );
            tokio::time::sleep(INTER_SWEEP_COOLDOWN).await;
            println!();
        } else {
            break;
        }
    }

    // ── Grand conclusion ──────────────────────────────────────────────────────
    println!("{}", "═".repeat(90).bright_cyan());
    println!("{}", "=== Final Report ===".bright_cyan().bold());
    println!();

    let n = measured_limits.len();
    if n == 0 {
        println!("  No sweeps completed.");
        return Ok(());
    }

    let avg_limit = measured_limits.iter().sum::<f64>() / n as f64;
    let min_limit = measured_limits
        .iter()
        .cloned()
        .fold(f64::INFINITY, f64::min);
    let max_limit = measured_limits
        .iter()
        .cloned()
        .fold(f64::NEG_INFINITY, f64::max);
    let safe_rate = (avg_limit * 0.8).max(0.5);

    println!("  Sweeps completed   : {n}");
    println!(
        "  Total elapsed      : {:.0}s",
        run_start.elapsed().as_secs_f64()
    );
    println!();

    for (idx, &limit) in measured_limits.iter().enumerate() {
        println!(
            "  Sweep #{:<3}  measured limit : {:.1} req/s",
            idx + 1,
            limit
        );
    }
    println!();

    println!(
        "  {:<22} {:.1} req/s",
        "Min measured limit :".bold(),
        min_limit
    );
    println!(
        "  {:<22} {:.1} req/s",
        "Max measured limit :".bold(),
        max_limit
    );
    println!(
        "  {:<22} {:.1} req/s",
        "Avg measured limit :".bold().bright_white(),
        avg_limit
    );
    println!();
    println!(
        "  {}  {:.1} req/s  (avg × 0.8)",
        "Recommended safe rate :".bold().bright_green(),
        safe_rate,
    );
    println!();
    println!("{}", "Done.".bright_green());
    Ok(())
}
