use std::convert::Infallible;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;

use apollo_blocking::{CancellationToken, ServiceBuilderExt};
use apollo_qos::concurrency_limit::ConcurrencyLimitLayer;
use futures::future::join_all;
use tower::{Service, ServiceBuilder, ServiceExt};

/// Dropping the adapter future cancels the child token.
#[tokio::test]
async fn dropped_future_cancels_blocking_thread() {
    let started = Arc::new(AtomicUsize::new(0));
    let observed = Arc::new(AtomicUsize::new(0));

    let svc = {
        let started = Arc::clone(&started);
        let observed = Arc::clone(&observed);
        ServiceBuilder::new().blocking_service(
            "drop_test",
            None,
            move |_req: (), ct: CancellationToken| -> Result<(), Infallible> {
                started.fetch_add(1, Ordering::SeqCst);
                while !ct.is_cancelled() {
                    std::thread::sleep(Duration::from_millis(1));
                }
                observed.fetch_add(1, Ordering::SeqCst);
                Ok(())
            },
        )
    };

    let mut svc = svc;
    let future = svc.call(());

    while started.load(Ordering::SeqCst) == 0 {
        tokio::task::yield_now().await;
    }

    drop(future);

    // Wait for the blocking thread to exit cooperatively.
    for _ in 0..1000 {
        if observed.load(Ordering::SeqCst) == 1 {
            return;
        }
        tokio::time::sleep(Duration::from_millis(10)).await;
    }
    panic!("blocking thread did not observe cancellation after future drop");
}

/// A parent token cancellation flows into the child observed by blocking code.
#[tokio::test]
async fn parent_token_cancellation_propagates_to_in_flight_task() {
    let parent = CancellationToken::new();
    let observed = Arc::new(AtomicUsize::new(0));

    let svc = {
        let observed = Arc::clone(&observed);
        ServiceBuilder::new().blocking_service(
            "parent_token",
            Some(parent.clone()),
            move |_req: (), ct: CancellationToken| -> Result<(), Infallible> {
                while !ct.is_cancelled() {
                    std::thread::sleep(Duration::from_millis(1));
                }
                observed.fetch_add(1, Ordering::SeqCst);
                Ok(())
            },
        )
    };

    let task = tokio::spawn(svc.oneshot(()));

    tokio::time::sleep(Duration::from_millis(20)).await;
    parent.cancel();

    task.await.unwrap().unwrap();
    assert_eq!(observed.load(Ordering::SeqCst), 1);
}

/// Sibling tasks under a shared parent token are isolated from each other's drops. Each call
/// must derive an independent child token from the parent, so dropping one task's future cancels
/// only that task and leaves concurrent siblings — and the parent itself — untouched.
#[tokio::test]
async fn dropping_one_future_does_not_cancel_sibling_under_shared_parent() {
    let parent = CancellationToken::new();
    let a_started = Arc::new(AtomicBool::new(false));
    let a_cancelled = Arc::new(AtomicBool::new(false));
    let b_started = Arc::new(AtomicBool::new(false));
    let b_cancelled = Arc::new(AtomicBool::new(false));
    let release_b = Arc::new(AtomicBool::new(false));

    let mut svc_a = {
        let started = Arc::clone(&a_started);
        let cancelled = Arc::clone(&a_cancelled);
        ServiceBuilder::new().blocking_service(
            "task_a",
            Some(parent.clone()),
            move |_req: (), ct: CancellationToken| -> Result<(), Infallible> {
                started.store(true, Ordering::SeqCst);
                while !ct.is_cancelled() {
                    std::thread::sleep(Duration::from_millis(1));
                }
                cancelled.store(true, Ordering::SeqCst);
                Ok(())
            },
        )
    };

    let svc_b = {
        let started = Arc::clone(&b_started);
        let cancelled = Arc::clone(&b_cancelled);
        let release = Arc::clone(&release_b);
        ServiceBuilder::new().blocking_service(
            "task_b",
            Some(parent.clone()),
            move |_req: (), ct: CancellationToken| -> Result<(), Infallible> {
                started.store(true, Ordering::SeqCst);
                while !release.load(Ordering::SeqCst) {
                    if ct.is_cancelled() {
                        cancelled.store(true, Ordering::SeqCst);
                        return Ok(());
                    }
                    std::thread::sleep(Duration::from_millis(1));
                }
                Ok(())
            },
        )
    };

    // A's future is held so we can drop it explicitly. B is awaited via a spawned task so it
    // can run to completion independently.
    let a_future = svc_a.call(());
    let b_handle = tokio::spawn(svc_b.oneshot(()));

    while !a_started.load(Ordering::SeqCst) || !b_started.load(Ordering::SeqCst) {
        tokio::task::yield_now().await;
    }

    drop(a_future);

    for _ in 0..1000 {
        if a_cancelled.load(Ordering::SeqCst) {
            break;
        }
        tokio::time::sleep(Duration::from_millis(10)).await;
    }
    assert!(
        a_cancelled.load(Ordering::SeqCst),
        "A did not observe cancellation after its future was dropped"
    );

    // Give any erroneous cross-task cancellation a moment to surface before asserting that
    // B and the parent are unaffected.
    tokio::time::sleep(Duration::from_millis(50)).await;
    assert!(
        !b_cancelled.load(Ordering::SeqCst),
        "B observed cancellation after sibling A was dropped — child tokens are not isolated"
    );
    assert!(
        !parent.is_cancelled(),
        "parent token was cancelled by a sibling's drop guard"
    );

    release_b.store(true, Ordering::SeqCst);
    b_handle.await.unwrap().unwrap();
    assert!(
        !b_cancelled.load(Ordering::SeqCst),
        "B observed cancellation during its run — child tokens are not isolated"
    );
}

/// A panic inside the blocking task propagates to the awaiting task with its payload intact.
#[tokio::test]
async fn panic_propagates_to_awaiting_task() {
    let svc = ServiceBuilder::new().blocking_service(
        "panic_test",
        None,
        |_req: (), _ct: CancellationToken| -> Result<(), Infallible> {
            panic!("intentional panic for testing");
        },
    );

    // Isolate the panic on a spawned task so it doesn't kill the test runner. The JoinHandle
    // reports the panic instead of an error variant — the panic flowed through the future.
    let join = tokio::spawn(svc.oneshot(())).await;

    let join_err = join.expect_err("blocking task panic should propagate to the awaiting task");
    assert!(join_err.is_panic());

    // Payload preserved across the resume_unwind chain.
    let payload = join_err.into_panic();
    let message = payload
        .downcast_ref::<&'static str>()
        .expect("panic!(\"literal\") payload is &'static str");
    assert_eq!(*message, "intentional panic for testing");
}

/// `apollo_qos::concurrency_limit::ConcurrencyLimitLayer::new(N)` above the adapter caps the
/// number of concurrent blocking tasks at N. The apollo-qos layer load-sheds excess requests
/// with a 503-style error rather than submitting them to the blocking pool, so concurrent
/// in-flight blocking work never exceeds the configured limit.
#[tokio::test]
async fn concurrency_limit_caps_in_flight_blocking_tasks() {
    const LIMIT: usize = 4;
    let in_flight = Arc::new(AtomicUsize::new(0));
    let max_observed = Arc::new(AtomicUsize::new(0));

    let svc = {
        let in_flight = Arc::clone(&in_flight);
        let max_observed = Arc::clone(&max_observed);
        ServiceBuilder::new()
            .layer(ConcurrencyLimitLayer::new(
                "concurrency_test",
                NonZeroUsize::new(LIMIT).unwrap(),
            ))
            .blocking_service(
                "concurrency_test",
                None,
                move |_req: (), _ct: CancellationToken| -> Result<(), Infallible> {
                    let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
                    max_observed.fetch_max(now, Ordering::SeqCst);
                    // Hold the slot long enough for the limit to bind on overlapping callers.
                    std::thread::sleep(Duration::from_millis(50));
                    in_flight.fetch_sub(1, Ordering::SeqCst);
                    Ok(())
                },
            )
    };

    // Burst more requests than the limit can admit. apollo-qos load-sheds the excess.
    let total = LIMIT * 4;
    let calls: Vec<_> = (0..total).map(|_| svc.clone().oneshot(())).collect();
    let results = join_all(calls).await;

    let successes = results.iter().filter(|r| r.is_ok()).count();
    let max = max_observed.load(Ordering::SeqCst);

    assert!(
        max <= LIMIT,
        "expected concurrent blocking tasks to stay within {LIMIT}, observed {max}"
    );
    assert!(successes > 0, "no requests succeeded — test setup broken");
    assert!(
        successes < total,
        "no requests were load-shed — concurrency limit was not exercised"
    );
}

/// Adapter `poll_ready` itself never refuses — concurrency limiting is the outer layer's job.
#[tokio::test]
async fn adapter_poll_ready_is_always_ready() {
    let svc = ServiceBuilder::new().blocking_service(
        "ready_test",
        None,
        |_req: (), _ct: CancellationToken| -> Result<(), Infallible> { Ok(()) },
    );

    // `ready_oneshot` would loop forever if poll_ready returned Pending; reaching the call
    // below proves it returned Ready immediately.
    let _ = svc.oneshot(()).await;
}
