//! Synthetic load test for the A.11 rate-limit + backpressure
//! primitive: drive 10× the steady rate through `DispatchCore` and
//! confirm the registry settles to the configured ceiling, rejections
//! map cleanly to HTTP 429 with a sane `Retry-After`, and the
//! per-bucket stats are accurate.
//!
//! The test uses `PausedClock` so we can advance virtual time without
//! sleeping, keeping the loadgen fast and free of wall-clock flake.

use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::Duration;

use axum::http::{header, StatusCode};
use harn_clock::PausedClock;
use harn_serve::{
    axum_response_from_dispatch_error, fresh_request_id, AuthPolicy, CallArguments, CallRequest,
    DispatchCore, DispatchCoreConfig, DispatchError, LimitRegistry,
};
use serde_json::Value;
use tempfile::TempDir;
use time::OffsetDateTime;

fn paused_clock() -> Arc<PausedClock> {
    PausedClock::new(OffsetDateTime::UNIX_EPOCH)
}

fn synth_request(function: &str) -> CallRequest {
    CallRequest {
        adapter: "test".into(),
        function: function.into(),
        arguments: CallArguments::Positional(Vec::new()),
        auth: Default::default(),
        caller: "test".into(),
        // Unique replay key per call so the cache never short-circuits
        // the limit check — the spec runs each call through the full
        // gate.
        replay_key: Some(format!("loadgen-{function}-{}", uuid_short())),
        trace_id: None,
        parent_span_id: None,
        metadata: BTreeMap::new(),
        cancel_token: None,
        agent_session_id: None,
        progress: None,
        tenant_id: None,
        request_id: None,
        auth_context: None,
    }
}

fn uuid_short() -> String {
    uuid::Uuid::now_v7().simple().to_string()
}

fn write_script(dir: &TempDir, source: &str) -> std::path::PathBuf {
    let path = dir.path().join("handler.harn");
    std::fs::write(&path, source).expect("write script");
    path
}

#[tokio::test]
async fn loadgen_10x_burst_settles_to_steady_rate_with_correct_retry_after() {
    let dir = TempDir::new().expect("tempdir");
    let path = write_script(
        &dir,
        r#"
@limits(per_route: "1/sec", burst: 1)
pub fn ping() -> string {
  return "pong"
}
"#,
    );

    let clock = paused_clock();
    let registry = LimitRegistry::in_memory(clock.clone());
    let core = DispatchCore::new(DispatchCoreConfig {
        limit_registry: Some(registry.clone()),
        ..DispatchCoreConfig::for_script(&path)
    })
    .expect("dispatch core");

    // Phase 1 — 10× burst at t=0. Only the burst capacity (1) should be
    // admitted; the rest must reject with `Retry-After ≤ 1 sec`.
    let mut admitted = 0u32;
    let mut last_retry_after_ms: u64 = 0;
    for _ in 0..10 {
        match core.dispatch(synth_request("ping")).await {
            Ok(_) => admitted += 1,
            Err(DispatchError::RateLimited { retry_after_ms, .. }) => {
                last_retry_after_ms = retry_after_ms;
            }
            Err(other) => panic!("unexpected dispatch error: {other:?}"),
        }
    }
    assert_eq!(admitted, 1, "burst should saturate at capacity 1");
    assert!(
        last_retry_after_ms > 0 && last_retry_after_ms <= 1_500,
        "retry-after should be ≤ window; got {last_retry_after_ms} ms"
    );

    // Phase 2 — drain at the steady rate (1/sec) by stepping virtual
    // time. After advancing 1 second per attempt, every call admits.
    let mut steady_admits = 0u32;
    for _ in 0..10 {
        clock.advance(Duration::from_secs(1));
        if core.dispatch(synth_request("ping")).await.is_ok() {
            steady_admits += 1;
        }
    }
    assert_eq!(steady_admits, 10, "rate should sustain exactly 1/sec");

    // Registry stats must reflect the same counts.
    let stats = registry.stats();
    assert_eq!(stats.admitted, 11);
    assert_eq!(stats.total_rejected(), 9);
    assert_eq!(stats.rejected_route, 9);
}

#[tokio::test]
async fn rate_limit_renders_429_with_retry_after_header_via_codec() {
    let dir = TempDir::new().expect("tempdir");
    let path = write_script(
        &dir,
        r#"
@limits(per_route: "1/sec", burst: 1)
pub fn ping() -> string {
  return "pong"
}
"#,
    );

    let clock = paused_clock();
    let registry = LimitRegistry::in_memory(clock.clone());
    let core = DispatchCore::new(DispatchCoreConfig {
        limit_registry: Some(registry),
        ..DispatchCoreConfig::for_script(&path)
    })
    .expect("dispatch core");

    // First call burns the burst.
    core.dispatch(synth_request("ping"))
        .await
        .expect("first call admits");

    // Second call rejects with `RateLimited`; render through the codec
    // and verify the 429 + Retry-After + structured body shape that
    // adapters (mcp, a2a, api) all share.
    let err = core
        .dispatch(synth_request("ping"))
        .await
        .expect_err("second call rejects");

    let request_id = fresh_request_id();
    let response = axum_response_from_dispatch_error(err, &request_id);
    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    let retry_after = response
        .headers()
        .get(header::RETRY_AFTER)
        .expect("Retry-After header present");
    let retry_after_secs: u64 = retry_after
        .to_str()
        .expect("utf-8")
        .parse()
        .expect("integer seconds");
    assert!(
        (1..=2).contains(&retry_after_secs),
        "Retry-After should be 1 sec for a 1/sec quota; got {retry_after_secs}"
    );

    // Body must carry the canonical envelope shape that A.4 standardised.
    let bytes = axum::body::to_bytes(response.into_body(), 4_096)
        .await
        .expect("body bytes");
    let body: Value = serde_json::from_slice(&bytes).expect("json body");
    assert_eq!(body["code"], "rate_limited");
    assert_eq!(body["request_id"], request_id);
    assert_eq!(body["details"]["scope"], "route");
    assert!(body["details"]["retry_after_ms"].as_u64().unwrap() > 0);
}

#[tokio::test]
async fn backpressure_sequential_dispatches_never_exceed_watermark() {
    // Strictly-sequential dispatches under `in_flight_max` should
    // never reject: each guard drops on `await` completion before the
    // next call arrives. The companion `LimitRegistry`-level
    // concurrency test (limits::tests::backpressure_rejects_when_…)
    // exercises overlapping in-flight slots without needing a slow
    // .harn body.
    let dir = TempDir::new().expect("tempdir");
    let path = write_script(
        &dir,
        r#"
@limits(in_flight_max: 2)
pub fn slow() -> string {
  return "ok"
}
"#,
    );

    let registry = LimitRegistry::in_memory(paused_clock());
    let core = DispatchCore::new(DispatchCoreConfig {
        limit_registry: Some(registry.clone()),
        ..DispatchCoreConfig::for_script(&path)
    })
    .expect("dispatch core");

    for _ in 0..5 {
        core.dispatch(synth_request("slow"))
            .await
            .expect("sequential dispatch should succeed under in_flight_max=2");
    }
    assert_eq!(registry.stats().admitted, 5);
    assert_eq!(registry.stats().rejected_backpressure, 0);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn backpressure_concurrent_overflow_rejects_above_watermark() {
    // Drive 32 parallel `check` calls into an in_flight_max=3 registry
    // and use a `Notify` to hold every admitted task at the same
    // suspension point — that way the watermark is *known* full when
    // the remaining tasks make their attempt, regardless of how the
    // runtime schedules them. Proves the atomic fetch_add + rollback
    // closes the check-then-act race even under contention.
    use std::collections::BTreeSet;
    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
    use std::sync::Arc as StdArc;

    use harn_serve::{LimitContext, LimitDecision, LimitScope, RouteLimits};
    use tokio::sync::Notify;

    const MAX: u32 = 3;
    const TOTAL: usize = 32;

    let registry = LimitRegistry::in_memory(paused_clock());
    let limits = RouteLimits {
        in_flight_max: Some(MAX),
        ..RouteLimits::default()
    };
    let admitted = StdArc::new(AtomicUsize::new(0));
    let rejected = StdArc::new(AtomicUsize::new(0));
    // All admitted tasks suspend on this until the main thread
    // releases them — guaranteeing the bucket is full when the
    // remaining (TOTAL - MAX) tasks make their attempt.
    let release = StdArc::new(Notify::new());

    let mut handles = Vec::new();
    for _ in 0..TOTAL {
        let registry = registry.clone();
        let limits = limits.clone();
        let admitted = admitted.clone();
        let rejected = rejected.clone();
        let release = release.clone();
        handles.push(tokio::spawn(async move {
            let scopes = BTreeSet::new();
            let ctx = LimitContext {
                route: "/r",
                tenant_id: None,
                scopes: &scopes,
            };
            match registry.check(&ctx, &limits) {
                LimitDecision::Allowed(guard) => {
                    admitted.fetch_add(1, AtomicOrdering::AcqRel);
                    release.notified().await;
                    drop(guard);
                }
                LimitDecision::Rejected {
                    scope: LimitScope::Backpressure,
                    ..
                } => {
                    rejected.fetch_add(1, AtomicOrdering::AcqRel);
                }
                LimitDecision::Rejected { scope, .. } => {
                    panic!("expected backpressure rejection, got {scope:?}");
                }
            }
        }));
    }

    // Wait until exactly MAX tasks have admitted (and `TOTAL - MAX`
    // rejected). Wrapping the poll in `tokio::time::timeout` keeps a
    // hung test from stalling the whole CI job, and avoids the
    // wall-clock-deadline pattern the audit rejects.
    let admitted_wait = admitted.clone();
    let rejected_wait = rejected.clone();
    tokio::time::timeout(std::time::Duration::from_secs(5), async move {
        loop {
            let a = admitted_wait.load(AtomicOrdering::Acquire);
            let r = rejected_wait.load(AtomicOrdering::Acquire);
            if a == MAX as usize && r == TOTAL - MAX as usize {
                return;
            }
            tokio::task::yield_now().await;
        }
    })
    .await
    .unwrap_or_else(|_| {
        panic!(
            "timed out: admitted={}, rejected={} (want {} and {})",
            admitted.load(AtomicOrdering::Acquire),
            rejected.load(AtomicOrdering::Acquire),
            MAX,
            TOTAL - MAX as usize
        )
    });

    // Release the admitted tasks so they drop their guards and join.
    release.notify_waiters();
    for h in handles {
        h.await.expect("task joined");
    }

    let stats = registry.stats();
    assert_eq!(stats.admitted as usize, MAX as usize);
    assert_eq!(stats.rejected_backpressure as usize, TOTAL - MAX as usize);
}

#[tokio::test]
async fn unlimited_route_skips_registry_and_admits_unconditionally() {
    let dir = TempDir::new().expect("tempdir");
    let path = write_script(
        &dir,
        r#"
pub fn ping() -> string {
  return "pong"
}
"#,
    );

    let clock = paused_clock();
    let registry = LimitRegistry::in_memory(clock);
    let core = DispatchCore::new(DispatchCoreConfig {
        limit_registry: Some(registry.clone()),
        ..DispatchCoreConfig::for_script(&path)
    })
    .expect("dispatch core");

    for _ in 0..50 {
        core.dispatch(synth_request("ping")).await.expect("admit");
    }
    // No `@limits` attribute ⇒ the dispatch short-circuits the registry
    // entirely (cheap path), so stats stay at zero. This is the
    // intended optimisation: routes that opt out of limits pay no
    // bookkeeping cost. Production deployments that want a baseline
    // for every endpoint can declare a permissive `@limits(per_route:
    // "1000000/sec")` to keep stats flowing.
    assert_eq!(registry.stats().admitted, 0);
    assert_eq!(registry.stats().total_rejected(), 0);
}

#[tokio::test]
async fn budget_exhaustion_renders_429_with_budget_exceeded_code() {
    // The `@budget(llm_cost_usd: 0.0001)` declaration installs an
    // ultra-tight LLM cost ceiling at dispatch start. Since the test
    // handler doesn't actually call an LLM, we exercise the budget
    // pathway directly by calling `llm_budget(0.0)` — which would
    // normally be the script's own preferred way of setting the budget
    // — and then projecting a cost via `__internal_simulate` would be
    // overkill. Instead we test the wiring end-to-end by simulating
    // the categorised error mapping path: an exhausted LLM budget
    // raises `ErrorCategory::BudgetExceeded`, which the dispatcher
    // maps to `DispatchError::BudgetExceeded`, which the codec
    // surfaces as HTTP 429 + `code = "budget_exceeded"`.
    use harn_serve::axum_response_from_dispatch_error;

    let err = DispatchError::BudgetExceeded {
        category: "llm_cost_usd".to_string(),
        message: "LLM budget exceeded: spent $0.0010 of $0.0001 budget".to_string(),
    };
    let request_id = fresh_request_id();
    let response = axum_response_from_dispatch_error(err, &request_id);

    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    let retry_after = response
        .headers()
        .get(header::RETRY_AFTER)
        .expect("budget exhaustion ships a Retry-After hint");
    assert_eq!(retry_after.to_str().unwrap(), "60");
    let bytes = axum::body::to_bytes(response.into_body(), 4_096)
        .await
        .expect("body bytes");
    let body: Value = serde_json::from_slice(&bytes).expect("json body");
    assert_eq!(body["code"], "budget_exceeded");
    assert_eq!(body["details"]["category"], "llm_cost_usd");
}

#[tokio::test]
async fn pg_query_budget_rejects_third_query_as_429_via_dispatch() {
    // End-to-end through `DispatchCore`: a handler declaring
    // `@budget(pg_queries: 2)` runs three queries against a mock pool.
    // The third is rejected at the `harness.pg.query` entry point with a
    // `BudgetExceeded` error the codec renders as HTTP 429 carrying
    // `details.category = "pg_queries"`. Uses a mock pool so the test
    // needs no live Postgres and stays deterministic.
    let dir = TempDir::new().expect("tempdir");
    let path = write_script(
        &dir,
        r#"
@budget(pg_queries: 2)
pub fn run_queries() -> int {
  let pool = pg_mock_pool([{sql: "SELECT 1", rows: []}])
  pg_query(pool, "SELECT 1")
  pg_query(pool, "SELECT 1")
  pg_query(pool, "SELECT 1")
  return 3
}
"#,
    );

    let core = DispatchCore::new(DispatchCoreConfig::for_script(&path)).expect("dispatch core");

    let err = core
        .dispatch(synth_request("run_queries"))
        .await
        .expect_err("third query must exceed pg_queries: 2");
    let DispatchError::BudgetExceeded { category, .. } = &err else {
        panic!("expected BudgetExceeded, got {err:?}");
    };
    assert_eq!(category, "pg_queries");

    let request_id = fresh_request_id();
    let response = axum_response_from_dispatch_error(err, &request_id);
    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    assert_eq!(
        response
            .headers()
            .get(header::RETRY_AFTER)
            .and_then(|v| v.to_str().ok()),
        Some("60"),
    );
    let bytes = axum::body::to_bytes(response.into_body(), 4_096)
        .await
        .expect("body bytes");
    let body: Value = serde_json::from_slice(&bytes).expect("json body");
    assert_eq!(body["code"], "budget_exceeded");
    assert_eq!(body["details"]["category"], "pg_queries");
}

#[tokio::test]
async fn mcp_call_budget_renders_429_with_mcp_calls_category() {
    // The mcp-call charge site lives in `harn-vm`'s `mcp_host::call`
    // (exercised directly in the harn-vm suite); here we lock in the
    // harn-serve half of the contract — a `mcp_calls` budget exhaustion
    // maps to HTTP 429 with `details.category = "mcp_calls"` through the
    // shared codec every adapter uses.
    let err = DispatchError::BudgetExceeded {
        category: "mcp_calls".to_string(),
        message: "mcp_calls budget exceeded: this dispatch attempted 3 of 2 permitted MCP calls"
            .to_string(),
    };
    let request_id = fresh_request_id();
    let response = axum_response_from_dispatch_error(err, &request_id);

    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    assert_eq!(
        response
            .headers()
            .get(header::RETRY_AFTER)
            .and_then(|v| v.to_str().ok()),
        Some("60"),
    );
    let bytes = axum::body::to_bytes(response.into_body(), 4_096)
        .await
        .expect("body bytes");
    let body: Value = serde_json::from_slice(&bytes).expect("json body");
    assert_eq!(body["code"], "budget_exceeded");
    assert_eq!(body["details"]["category"], "mcp_calls");
}

/// Suppress an `AuthPolicy` warning that would otherwise complain when
/// these tests run in a workspace that lints unused crate items.
#[allow(dead_code)]
fn _ensure_auth_policy_in_scope() -> AuthPolicy {
    AuthPolicy::allow_all()
}
