//! Per-upstream request budget — a shared token bucket.
//!
//! Lifted from the card feeder's `tmdb_budget.rs` (its `TmdbBudget` half; the
//! per-indexer registry has no counterpart here) and renamed, because the
//! mechanism is upstream-agnostic and the music feeder needs three independent
//! instances rather than one global.
//!
//! **Why three, not one.** MusicBrainz's 1 req/s is a hard published ceiling
//! enforced per source IP; ListenBrainz and the Internet Archive are far more
//! generous. A single shared bucket would either throttle all three to 1/s or
//! let MusicBrainz calls burst past a limit that gets the whole box blocked.
//!
//! Design notes, unchanged from the original:
//! - **Lazy refill.** Tokens accrue on each `acquire` from elapsed wall-clock,
//!   so there is no background timer task and no idle cost.
//! - **Cancel-safe.** The token is decremented synchronously under the lock
//!   right before `acquire` returns [`Lease::Granted`]; there is no `.await`
//!   between the decision to grant and the return, so a dropped `acquire`
//!   future (consumer disconnected) never spends a token.
//! - **Global pause on throttle.** One upstream 429/503 freezes *all* grants
//!   on that bucket until the Retry-After window elapses, via
//!   [`RateBudget::note_throttled`].
//!
//! The critical sections hold a plain `std::sync::Mutex` (no awaits inside),
//! while the waiting happens outside the lock on a `tokio::time::sleep` raced
//! against a `tokio::sync::Notify`.

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use tokio::sync::Notify;

/// Outcome of a [`RateBudget::acquire`] call.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Lease {
    /// A token was granted; the caller may make one upstream call.
    Granted,
    /// The wait deadline elapsed before a token became available. The caller
    /// degrades best-effort — it does **not** call the upstream anyway.
    DeadlineExceeded,
}

struct BucketState {
    tokens: f64,
    capacity: f64,
    refill_per_sec: f64,
    last_refill: Instant,
    /// When set and in the future, all grants are frozen.
    paused_until: Option<Instant>,
}

impl BucketState {
    /// Accrue tokens for elapsed time since the last refill. Called under the
    /// lock before inspecting `tokens`.
    fn refill(&mut self, now: Instant) {
        let elapsed = now.saturating_duration_since(self.last_refill).as_secs_f64();
        if elapsed > 0.0 {
            self.tokens = (self.tokens + elapsed * self.refill_per_sec).min(self.capacity);
            self.last_refill = now;
        }
    }
}

/// Shared token bucket gating one upstream's API calls across the process.
pub struct RateBudget {
    state: Mutex<BucketState>,
    /// Woken whenever a grant *might* now be possible (a token accrued, or a
    /// pause was cleared). Waiters re-check under the lock after waking; a
    /// missed wake only costs an extra bounded sleep.
    wake: Notify,
}

impl RateBudget {
    /// Build a budget sustaining `refill_per_sec` grants/second with a burst
    /// ceiling of `capacity` tokens. The bucket starts full, so the first
    /// burst is served immediately.
    pub fn new(refill_per_sec: f64, capacity: f64) -> Arc<Self> {
        Arc::new(Self {
            state: Mutex::new(BucketState {
                tokens: capacity,
                capacity,
                refill_per_sec,
                last_refill: Instant::now(),
                paused_until: None,
            }),
            wake: Notify::new(),
        })
    }

    /// Wait for one token, giving up after `deadline`. Cancel-safe: dropping
    /// the returned future before it resolves [`Lease::Granted`] spends no
    /// token.
    pub async fn acquire(&self, deadline: Duration) -> Lease {
        let give_up_at = Instant::now() + deadline;
        loop {
            // Decide-or-compute-wait, all under the lock (no awaits inside).
            let sleep_for = {
                let mut st = self.state.lock().expect("rate budget mutex poisoned");
                let now = Instant::now();
                if let Some(until) = st.paused_until {
                    if now >= until {
                        st.paused_until = None;
                    }
                }
                let paused = st.paused_until;
                if paused.is_none() {
                    st.refill(now);
                    if st.tokens >= 1.0 {
                        // Atomic grant: decrement and return with no await between.
                        st.tokens -= 1.0;
                        return Lease::Granted;
                    }
                }
                if now >= give_up_at {
                    return Lease::DeadlineExceeded;
                }
                let until_deadline = give_up_at.saturating_duration_since(now);
                let wait = match paused {
                    // Frozen: wake when the pause is scheduled to lift.
                    Some(until) => until.saturating_duration_since(now),
                    // Throttled: wake when the next whole token should accrue.
                    None => {
                        let needed = 1.0 - st.tokens;
                        let secs = if st.refill_per_sec > 0.0 {
                            needed / st.refill_per_sec
                        } else {
                            until_deadline.as_secs_f64()
                        };
                        Duration::from_secs_f64(secs.max(0.0))
                    }
                };
                wait.min(until_deadline)
            };
            tokio::select! {
                _ = tokio::time::sleep(sleep_for) => {}
                _ = self.wake.notified() => {}
            }
            if Instant::now() >= give_up_at {
                return Lease::DeadlineExceeded;
            }
        }
    }

    /// Feed back an upstream throttle response (429, or MusicBrainz's 503):
    /// pause all grants on this bucket until `now + retry_after`. Extends an
    /// existing pause but never shortens it.
    pub fn note_throttled(&self, retry_after: Duration) {
        let until = Instant::now() + retry_after;
        {
            let mut st = self.state.lock().expect("rate budget mutex poisoned");
            st.paused_until = Some(match st.paused_until {
                Some(existing) if existing > until => existing,
                _ => until,
            });
        }
        // Wake waiters so they recompute their sleep against the new pause.
        self.wake.notify_waiters();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test(start_paused = true)]
    async fn grants_up_to_capacity_immediately() {
        let b = RateBudget::new(3.0, 3.0);
        for _ in 0..3 {
            assert_eq!(b.acquire(Duration::from_secs(10)).await, Lease::Granted);
        }
    }

    #[tokio::test(start_paused = true)]
    async fn next_token_waits_for_refill() {
        let b = RateBudget::new(2.0, 2.0); // 2/s sustained
        for _ in 0..2 {
            assert_eq!(b.acquire(Duration::from_secs(10)).await, Lease::Granted);
        }
        // Bucket empty; next grant must wait ~0.5 s for one token at 2/s.
        let start = Instant::now();
        assert_eq!(b.acquire(Duration::from_secs(10)).await, Lease::Granted);
        assert!(
            start.elapsed() >= Duration::from_millis(450),
            "expected to wait ~0.5s for refill, waited {:?}",
            start.elapsed()
        );
    }

    /// The MusicBrainz shape specifically: 1/s with a burst of 2 must not let
    /// three calls through inside one second. This is the property that keeps
    /// the box off their block-list.
    #[tokio::test(start_paused = true)]
    async fn musicbrainz_shape_holds_one_per_second() {
        let b = RateBudget::new(1.0, 1.0);
        let start = Instant::now();
        for _ in 0..4 {
            assert_eq!(b.acquire(Duration::from_secs(30)).await, Lease::Granted);
        }
        // 4 grants from a 2-token bucket refilling at 1/s ⇒ at least 2 s.
        assert!(
            start.elapsed() >= Duration::from_millis(1_900),
            "4 calls at 1/s (burst 2) must take ~2s, took {:?}",
            start.elapsed()
        );
    }

    #[tokio::test(start_paused = true)]
    async fn deadline_exceeded_when_drained() {
        let b = RateBudget::new(0.001, 1.0); // effectively no refill within the test
        assert_eq!(b.acquire(Duration::from_secs(1)).await, Lease::Granted);
        assert_eq!(
            b.acquire(Duration::from_millis(200)).await,
            Lease::DeadlineExceeded
        );
    }

    #[tokio::test(start_paused = true)]
    async fn note_throttled_pauses_grants_even_with_tokens_available() {
        let b = RateBudget::new(100.0, 100.0);
        b.note_throttled(Duration::from_secs(2));
        assert_eq!(
            b.acquire(Duration::from_millis(500)).await,
            Lease::DeadlineExceeded
        );
        // After the pause window, grants resume.
        assert_eq!(b.acquire(Duration::from_secs(5)).await, Lease::Granted);
    }

    /// A pause is extended, never shortened — otherwise a late short retry
    /// would cancel an earlier long one and walk straight back into the block.
    #[tokio::test(start_paused = true)]
    async fn a_shorter_retry_after_does_not_shorten_an_existing_pause() {
        let b = RateBudget::new(100.0, 100.0);
        b.note_throttled(Duration::from_secs(10));
        b.note_throttled(Duration::from_secs(1));
        assert_eq!(
            b.acquire(Duration::from_secs(3)).await,
            Lease::DeadlineExceeded
        );
    }
}