//! Per-tenant LLM governor (Phase B B4).
//!
//! Drives `Governor::acquire_llm_scoped` on `TokenBucketGovernor` across the
//! three contracts the inbound-MCP budget gate depends on:
//!
//! - `BudgetScope::Provider(_)` still routes the existing provider bucket
//!   (back-compat with `acquire_llm` + fail-closed on unknown).
//! - `BudgetScope::TenantProvider(_, _)` isolates tenants, applies a
//!   *bounded* default cap to unconfigured tenants, and honours
//!   `tenant_llm_limit` overrides.
//! - The default trait impl on a stub `Governor` falls back to the
//!   provider-scoped `acquire_llm`.

use async_trait::async_trait;
use klieo_bus_memory::MemoryBus;
use klieo_ops::governor::{Governor, GovernorError, Host, Permit, TokenBucketGovernor};
use klieo_ops::types::{BudgetScope, ProviderId, TenantId};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;

const ANTHROPIC: &str = "anthropic";
const TENANT_A: &str = "tenant_a";
const TENANT_B: &str = "tenant_b";
const TENANT_OVERRIDE: &str = "tenant_with_override";

const PROVIDER_CAP: u32 = 5;
const TENANT_DEFAULT_CAP: u32 = 2;
const TENANT_OVERRIDE_CAP: u32 = 4;
const SATURATION_PROBE_DEPTH: u32 = 8;
const EST_TOKENS: u32 = 100;

/// Drain a scope until it refuses. Returns the held permits + the error
/// the first refusal yielded so a caller can assert on both depth and
/// failure shape.
async fn drain_scope(
    g: &TokenBucketGovernor,
    scope: BudgetScope,
) -> (Vec<Permit>, GovernorError) {
    let mut held = Vec::new();
    for _ in 0..SATURATION_PROBE_DEPTH {
        match g.acquire_llm_scoped(scope.clone(), EST_TOKENS).await {
            Ok(p) => held.push(p),
            Err(e) => return (held, e),
        }
    }
    panic!(
        "scope {scope:?} never saturated within {SATURATION_PROBE_DEPTH} acquires; \
         bounded-default test would not be falsifiable"
    );
}

fn build_governor_with_tenant_default() -> TokenBucketGovernor {
    let bus = MemoryBus::new();
    TokenBucketGovernor::builder(bus.kv.clone())
        .llm_limit(ProviderId(ANTHROPIC.into()), PROVIDER_CAP)
        .tenant_llm_default_rps(TENANT_DEFAULT_CAP)
        .tenant_llm_limit(
            TenantId(TENANT_OVERRIDE.into()),
            ProviderId(ANTHROPIC.into()),
            TENANT_OVERRIDE_CAP,
        )
        .build()
}

#[tokio::test]
async fn tenant_buckets_are_isolated() {
    let g = build_governor_with_tenant_default();
    let provider = ProviderId(ANTHROPIC.into());

    let (held_a, _) = drain_scope(
        &g,
        BudgetScope::TenantProvider(TenantId(TENANT_A.into()), provider.clone()),
    )
    .await;

    // Tenant B must still acquire while tenant A is saturated.
    let permit_b = g
        .acquire_llm_scoped(
            BudgetScope::TenantProvider(TenantId(TENANT_B.into()), provider.clone()),
            EST_TOKENS,
        )
        .await
        .expect("tenant B is not saturated by tenant A's drain");

    drop(permit_b);
    drop(held_a);
}

#[tokio::test]
async fn unconfigured_tenant_uses_bounded_default() {
    let g = build_governor_with_tenant_default();
    let provider = ProviderId(ANTHROPIC.into());

    let (held, err) = drain_scope(
        &g,
        BudgetScope::TenantProvider(TenantId(TENANT_A.into()), provider),
    )
    .await;

    // Bounded means: depth equals the configured default cap exactly,
    // and the refusal is `Saturated`. Unbounded fallback would hit the
    // `SATURATION_PROBE_DEPTH` panic above.
    assert_eq!(
        held.len(),
        TENANT_DEFAULT_CAP as usize,
        "unconfigured tenant must saturate at tenant_llm_default_rps"
    );
    assert!(
        matches!(err, GovernorError::Saturated { .. } | GovernorError::TimedOut { .. }),
        "expected Saturated/TimedOut once default cap is drained, got {err:?}"
    );
}

#[tokio::test]
async fn tenant_llm_limit_override_is_honoured() {
    let g = build_governor_with_tenant_default();
    let provider = ProviderId(ANTHROPIC.into());

    let (held, _) = drain_scope(
        &g,
        BudgetScope::TenantProvider(TenantId(TENANT_OVERRIDE.into()), provider),
    )
    .await;

    assert_eq!(
        held.len(),
        TENANT_OVERRIDE_CAP as usize,
        "tenant_llm_limit override must replace the default cap"
    );
}

#[tokio::test]
async fn provider_scope_path_is_unchanged() {
    let g = build_governor_with_tenant_default();
    let provider = ProviderId(ANTHROPIC.into());

    // The Provider path keeps the legacy wait-with-30s-timeout
    // semantics: a saturated bucket back-pressures the caller instead
    // of failing fast. Observe the cap via the budget snapshot — a
    // bounded probe would either succeed forever (refiller tops up)
    // or hang (timeout) and prove nothing.
    let snapshot = g.budget(BudgetScope::Provider(provider.clone())).await;
    assert_eq!(
        snapshot.limit, i64::from(PROVIDER_CAP),
        "Provider-scoped bucket must keep its configured cap"
    );

    // The bare `acquire_llm` path must still fail-closed on unknown.
    let unknown = g
        .acquire_llm(ProviderId("not-registered".into()), EST_TOKENS)
        .await;
    let Err(err) = unknown else {
        panic!("unknown provider must fail-closed, not grant a permit");
    };
    assert!(
        matches!(err, GovernorError::Saturated { .. }),
        "unknown provider must fail-closed via Saturated, got {err:?}"
    );

    // And `acquire_llm_scoped(Provider(_))` must reach the same
    // bucket the legacy `acquire_llm` reaches: an unknown provider
    // surfaces the same `Saturated` shape.
    let scoped_unknown = g
        .acquire_llm_scoped(BudgetScope::Provider(ProviderId("not-registered".into())), EST_TOKENS)
        .await;
    let Err(err) = scoped_unknown else {
        panic!("unknown provider via scoped path must fail-closed");
    };
    assert!(matches!(err, GovernorError::Saturated { .. }));
}

#[tokio::test]
async fn other_scopes_yield_unavailable() {
    let g = build_governor_with_tenant_default();
    for scope in [
        BudgetScope::Global,
        BudgetScope::Tenant(TenantId(TENANT_A.into())),
        BudgetScope::Agent(klieo_ops::AgentId("ag".into())),
    ] {
        let res = g.acquire_llm_scoped(scope.clone(), EST_TOKENS).await;
        let Err(err) = res else {
            panic!(
                "scope {scope:?} must NOT grant a permit (would silently bypass the gate)"
            );
        };
        assert!(
            matches!(err, GovernorError::Unavailable(_)),
            "scope {scope:?} must surface Unavailable, got {err:?}"
        );
    }
}

/// A stub `Governor` that records every call. We use it to prove the
/// `acquire_llm_scoped` default impl on the trait delegates to
/// `acquire_llm` so legacy impls keep working with no source edits.
struct CountingGovernor {
    llm_calls: AtomicU32,
}

#[async_trait]
impl Governor for CountingGovernor {
    async fn acquire_llm(
        &self,
        _provider: ProviderId,
        _est_tokens: u32,
    ) -> Result<Permit, GovernorError> {
        self.llm_calls.fetch_add(1, Ordering::Relaxed);
        Ok(Permit::new(|| {}))
    }

    async fn acquire_egress(&self, _host: Host) -> Result<Permit, GovernorError> {
        Ok(Permit::new(|| {}))
    }

    async fn budget(
        &self,
        _scope: BudgetScope,
    ) -> klieo_ops::governor::BudgetSnapshot {
        // `BudgetSnapshot` is `#[non_exhaustive]`; use the public `..Default::default()`
        // shim so future fields are absorbed automatically.
        klieo_ops::governor::BudgetSnapshot::default()
    }

    async fn fence(
        &self,
        _scope: BudgetScope,
        _limit: klieo_ops::governor::BudgetLimit,
    ) -> Result<(), GovernorError> {
        Ok(())
    }
}

#[tokio::test]
async fn default_impl_delegates_to_acquire_llm() {
    let g: Arc<dyn Governor> = Arc::new(CountingGovernor {
        llm_calls: AtomicU32::new(0),
    });
    let provider = ProviderId(ANTHROPIC.into());

    // Provider variant + TenantProvider variant must both reach
    // `acquire_llm` on the legacy impl.
    let _ = g
        .acquire_llm_scoped(BudgetScope::Provider(provider.clone()), EST_TOKENS)
        .await
        .expect("Provider scope delegates");
    let _ = g
        .acquire_llm_scoped(
            BudgetScope::TenantProvider(TenantId(TENANT_A.into()), provider),
            EST_TOKENS,
        )
        .await
        .expect("TenantProvider scope delegates");

    // Cast back to assert the recorded counter without depending on
    // dyn-trait downcast.
    let counting = Arc::into_raw(g) as *const CountingGovernor;
    // SAFETY: the Arc was minted from a `CountingGovernor`; we round-trip
    // it to read the counter then rebuild the Arc so its refcount stays
    // honest and Drop runs once.
    let counter = unsafe {
        let count = (*counting).llm_calls.load(Ordering::Relaxed);
        Arc::<CountingGovernor>::from_raw(counting);
        count
    };
    assert_eq!(
        counter, 2,
        "default impl must call acquire_llm once per scoped call"
    );
}
