//! Polling loop and backoff schedule for `--wait-for-ready`.
//!
//! Why split this into a [`ReadyProbe`] trait plus a [`wait_for_ready`]
//! driver? The runner supplies an HTTP-backed probe that opens a
//! TCP connection and sends an `initialize` JSON-RPC frame. Pulling reqwest
//! into the readiness crate would force every consumer (Python FFI host,
//! cassette replayer) to link an HTTP stack just to wait for a server to
//! come up. The trait keeps the polling logic transport-free and unit
//! testable.

use std::thread;
use std::time::{Duration, Instant};

use crate::network::Layer;

/// Outcome of polling a server for readiness.
///
/// `Ready` is the success path. `Timeout` means the probe never returned
/// success within the budget; the CLI maps it to exit code 3. `FailFast`
/// covers errors the runner cannot retry past (401, protocol mismatch,
/// bad TLS cert) and is exit code 2 at the CLI layer.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReadinessOutcome {
    /// The probe returned success before the timeout expired.
    Ready,
    /// The budget elapsed before the probe returned success.
    Timeout,
    /// The probe reported a non-retryable error.
    FailFast,
}

/// Reason a probe asked the driver to stop retrying.
///
/// Carries the network [`Layer`] so the CLI can render the same colored
/// banner it uses for runtime errors and a message the user can act on.
#[derive(Debug, Clone)]
pub struct ReadinessFailFast {
    /// Layer the probe diagnosed (typically `Authentication` or
    /// `McpInitialize`).
    pub layer: Layer,
    /// Human-readable description; surfaced verbatim by the CLI.
    pub message: String,
}

/// Strategy abstraction for a single readiness check.
///
/// Implementations return `Ok(true)` once the server accepts `initialize`,
/// `Ok(false)` for transient failures the driver should retry past (DNS,
/// TCP, TLS, 5xx) and `Err(...)` for protocol-level errors that should
/// stop the loop immediately.
pub trait ReadyProbe: Send + Sync {
    /// Run one probe attempt. See the trait docs for the return contract.
    fn probe(&self) -> Result<bool, ReadinessFailFast>;
}

/// Backoff schedule used between probe attempts, in milliseconds.
///
/// The schedule mirrors the spec on : short retries while the
/// server is likely still starting (250ms, 500ms, 1s, 2s) then a 5s
/// cap that keeps probe rate sane during long preview-deploy waits.
const BACKOFF_MS: &[u64] = &[250, 500, 1000, 2000, 5000];

/// Return the backoff delay for the `attempt`-th retry (zero-indexed).
///
/// Capped at 5s; callers that hand in any value get a deterministic delay
/// rather than a panic, which matters for property tests that explore the
/// tail of the schedule.
pub fn backoff_for_attempt(attempt: usize) -> Duration {
    let idx = attempt.min(BACKOFF_MS.len() - 1);
    Duration::from_millis(BACKOFF_MS[idx])
}

/// Poll `probe` until it succeeds or `timeout` elapses.
///
/// Calls `tick` roughly once per elapsed second so the CLI can print an
/// "elapsed N s" progress line. The tick fires whenever the wall clock
/// crosses an N-second boundary since the call started; under heavy load a
/// tick may slip but the count over the full span stays accurate.
///
/// The loop sleeps in 50ms slices so it can hit timeouts and tick boundaries
/// at sub-second precision without blocking on a long sleep.
pub fn wait_for_ready<P, F>(probe: &P, timeout: Duration, mut tick: F) -> ReadinessOutcome
where
    P: ReadyProbe,
    F: FnMut(Duration),
{
    let start = Instant::now();
    let deadline = start + timeout;
    let mut attempt: usize = 0;
    let mut next_tick = start + Duration::from_secs(1);
    let slice = Duration::from_millis(50);

    loop {
        match probe.probe() {
            Ok(true) => return ReadinessOutcome::Ready,
            Err(_) => return ReadinessOutcome::FailFast,
            Ok(false) => {}
        }

        let delay = backoff_for_attempt(attempt);
        attempt += 1;
        let sleep_until = Instant::now() + delay;

        loop {
            let now = Instant::now();
            if now >= deadline {
                return ReadinessOutcome::Timeout;
            }
            while now >= next_tick {
                tick(next_tick.duration_since(start));
                next_tick += Duration::from_secs(1);
            }
            if now >= sleep_until {
                break;
            }
            let remaining_to_sleep = sleep_until.saturating_duration_since(now);
            let remaining_to_deadline = deadline.saturating_duration_since(now);
            let remaining_to_tick = next_tick.saturating_duration_since(now);
            let step = slice
                .min(remaining_to_sleep)
                .min(remaining_to_deadline)
                .min(remaining_to_tick);
            thread::sleep(step);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Mutex;

    struct ScriptedProbe {
        results: Mutex<Vec<Result<bool, ReadinessFailFast>>>,
        calls: AtomicUsize,
    }

    impl ScriptedProbe {
        fn new(results: Vec<Result<bool, ReadinessFailFast>>) -> Self {
            Self {
                results: Mutex::new(results),
                calls: AtomicUsize::new(0),
            }
        }

        fn always_false() -> Self {
            // Empty queue means probe() falls through to the default
            // `Ok(false)` branch on every call, simulating a server that
            // never comes up.
            Self {
                results: Mutex::new(Vec::new()),
                calls: AtomicUsize::new(0),
            }
        }
    }

    impl ReadyProbe for ScriptedProbe {
        fn probe(&self) -> Result<bool, ReadinessFailFast> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            let mut guard = self.results.lock().unwrap();
            if guard.is_empty() {
                // Default: keep returning false so the driver keeps polling.
                return Ok(false);
            }
            guard.remove(0)
        }
    }

    #[test]
    fn ready_on_first_attempt_returns_quickly() {
        let probe = ScriptedProbe::new(vec![Ok(true)]);
        let start = Instant::now();
        let outcome = wait_for_ready(&probe, Duration::from_secs(5), |_| {});
        assert_eq!(outcome, ReadinessOutcome::Ready);
        assert_eq!(probe.calls.load(Ordering::SeqCst), 1);
        assert!(
            start.elapsed() < Duration::from_millis(200),
            "expected immediate return, took {:?}",
            start.elapsed()
        );
    }

    #[test]
    fn ready_after_n_transient_failures() {
        let probe = ScriptedProbe::new(vec![Ok(false), Ok(false), Ok(true)]);
        let outcome = wait_for_ready(&probe, Duration::from_secs(5), |_| {});
        assert_eq!(outcome, ReadinessOutcome::Ready);
        assert_eq!(probe.calls.load(Ordering::SeqCst), 3);
    }

    #[test]
    fn always_false_probe_times_out() {
        let probe = ScriptedProbe::always_false();
        let start = Instant::now();
        let outcome = wait_for_ready(&probe, Duration::from_millis(600), |_| {});
        assert_eq!(outcome, ReadinessOutcome::Timeout);
        // Allow generous slack for slow CI runners.
        assert!(
            start.elapsed() < Duration::from_secs(2),
            "timeout overshot, took {:?}",
            start.elapsed()
        );
        assert!(probe.calls.load(Ordering::SeqCst) >= 1);
    }

    #[test]
    fn fail_fast_error_stops_immediately() {
        let probe = ScriptedProbe::new(vec![Err(ReadinessFailFast {
            layer: Layer::Authentication,
            message: "401 unauthorized".into(),
        })]);
        let outcome = wait_for_ready(&probe, Duration::from_secs(5), |_| {});
        assert_eq!(outcome, ReadinessOutcome::FailFast);
        assert_eq!(probe.calls.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn tick_callback_fires_once_per_second_across_two_seconds() {
        let probe = ScriptedProbe::always_false();
        let counter = AtomicUsize::new(0);
        let outcome = wait_for_ready(&probe, Duration::from_millis(2100), |_elapsed| {
            counter.fetch_add(1, Ordering::SeqCst);
        });
        assert_eq!(outcome, ReadinessOutcome::Timeout);
        let ticks = counter.load(Ordering::SeqCst);
        // The callback fires about once per second, so ~2 ticks across 2.1s.
        // A loaded CI runner delays the per-second ticks, so the count can drop
        // (only the 1s tick lands before the 2.1s timeout); that is not a bug.
        // Assert the periodic callback fires at all, with a generous ceiling
        // that still catches a runaway tick loop.
        assert!(
            (1..=10).contains(&ticks),
            "expected the per-second tick to fire 1 to 10 times across ~2.1s, got {ticks}"
        );
    }

    #[test]
    fn backoff_schedule_matches_spec() {
        let expected_ms = [250u64, 500, 1000, 2000, 5000, 5000, 5000];
        for (attempt, ms) in expected_ms.iter().enumerate() {
            assert_eq!(
                backoff_for_attempt(attempt),
                Duration::from_millis(*ms),
                "attempt {attempt}"
            );
        }
    }
}
