use crate::clock::SystemClock;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;

/// A timeout wrapper for futures that returns the provided error if the future
/// does not complete within the specified duration.
///
/// # Arguments:
/// - `clock`: The clock to use for the timeout.
/// - `duration`: The duration to wait for the future to complete.
/// - `error_fn`: Returns the error to use when the timeout expires.
/// - `future`: The future to timeout
///
/// # Returns:
/// - `Ok(T)`: If the future completes within the specified duration.
/// - `Err(Err)`: If the future does not complete within the specified duration.
pub async fn timeout<T, Err>(
    clock: Arc<dyn SystemClock>,
    duration: Duration,
    error_fn: impl FnOnce() -> Err,
    future: impl Future<Output = Result<T, Err>> + Send,
) -> Result<T, Err> {
    tokio::select! {
        biased;
        res = future => res,
        _ = clock.sleep(duration) => Err(error_fn())
    }
}

#[cfg(all(test, feature = "test-util"))]
mod tests {
    use super::timeout;
    use crate::clock::MockSystemClock;
    use crate::clock::SystemClock;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::Arc;
    use std::time::Duration;

    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    enum TestError {
        Timeout,
    }

    #[tokio::test]
    async fn test_timeout_completes_before_expiry() {
        // Given: a mock clock and a future that completes quickly
        let clock = Arc::new(MockSystemClock::new());

        // When: we execute a future with a timeout
        let completed_future = async { Ok::<_, TestError>(42) };
        let timeout_future = timeout(
            clock,
            Duration::from_millis(100),
            || TestError::Timeout,
            completed_future,
        );

        // Then: the future should complete successfully with the expected value
        let result = timeout_future.await;
        assert_eq!(result.unwrap(), 42);
    }

    #[tokio::test]
    async fn test_timeout_expires() {
        // Given: a mock clock and a future that will never complete
        let clock = Arc::new(MockSystemClock::new());
        let never_completes = std::future::pending::<Result<(), TestError>>();
        let timeout_duration = Duration::from_millis(100);

        // When: we execute the future with a timeout and advance the clock past the timeout duration
        let timeout_future = timeout(
            clock.clone(),
            timeout_duration,
            || TestError::Timeout,
            never_completes,
        );
        let done = Arc::new(AtomicBool::new(false));
        let this_done = done.clone();

        tokio::spawn(async move {
            while !this_done.load(Ordering::SeqCst) {
                clock.advance(Duration::from_millis(100)).await;
                // Yield or else the scheduler keeps picking this loop, which
                // the sleep task forever.
                tokio::task::yield_now().await;
            }
        });

        // Then: the future should complete with a timeout error
        let result = timeout_future.await;
        done.store(true, Ordering::SeqCst);
        assert_eq!(result, Err(TestError::Timeout));
    }

    #[tokio::test]
    async fn test_timeout_respects_biased_select() {
        // Given: a mock clock and two futures that complete simultaneously
        let clock = Arc::new(MockSystemClock::new());
        let completes_immediately = async { Ok::<_, TestError>(42) };

        // When: we execute the future with a timeout and both are ready immediately
        let timeout_future = timeout(
            clock,
            Duration::from_millis(100),
            || TestError::Timeout,
            completes_immediately,
        );

        // Then: because of the 'biased' select, the future should complete with the value
        // rather than timing out, even though both are ready
        let result = timeout_future.await;
        assert_eq!(result.unwrap(), 42);
    }
}
