// file: crates/ksp-logging-lib/tests/tokio_span.rs
// version: 1

//! Integration tests for KSP span instrumentation on a real Tokio executor.

const TEST_TARGET: &str = "ksp-logging-lib";

#[derive(Clone)]
struct CountingSubscriber {
    enters: std::sync::Arc<std::sync::atomic::AtomicU64>,
    exits: std::sync::Arc<std::sync::atomic::AtomicU64>,
    next_id: std::sync::Arc<std::sync::atomic::AtomicU64>,
}

impl CountingSubscriber {
    fn new(enters: std::sync::Arc<std::sync::atomic::AtomicU64>, exits: std::sync::Arc<std::sync::atomic::AtomicU64>) -> Self {
        return Self { enters, exits, next_id: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(1)) };
    }
}

impl tracing::Subscriber for CountingSubscriber {
    fn enabled(&self, _metadata: &tracing::Metadata<'_>) -> bool {
        return true;
    }

    fn new_span(&self, _span: &tracing::span::Attributes<'_>) -> tracing::span::Id {
        let id = self.next_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        return tracing::span::Id::from_u64(id);
    }

    fn record(&self, _span: &tracing::span::Id, _values: &tracing::span::Record<'_>) {
        return;
    }

    fn record_follows_from(&self, _span: &tracing::span::Id, _follows: &tracing::span::Id) {
        return;
    }

    fn event(&self, _event: &tracing::Event<'_>) {
        return;
    }

    fn enter(&self, _span: &tracing::span::Id) {
        self.enters.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        return;
    }

    fn exit(&self, _span: &tracing::span::Id) {
        self.exits.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        return;
    }
}

fn test_span(enters: std::sync::Arc<std::sync::atomic::AtomicU64>, exits: std::sync::Arc<std::sync::atomic::AtomicU64>) -> ksp_logging_lib::Span {
    let subscriber = CountingSubscriber::new(enters, exits);
    return tracing::subscriber::with_default(subscriber, || -> ksp_logging_lib::Span {
        return ksp_logging_lib::trace_span!(target: TEST_TARGET, "tokio_runtime_span", domain = "logging", executor = "tokio");
    });
}

#[tokio::test(flavor = "current_thread")]
async fn instrumented_span_reenters_across_real_tokio_suspensions() {
    let enters = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let exits = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let span = test_span(std::sync::Arc::clone(&enters), std::sync::Arc::clone(&exits));
    let observed_enters = std::sync::Arc::clone(&enters);
    let future = ksp_logging_lib::instrument(span, async move {
        assert!(observed_enters.load(std::sync::atomic::Ordering::Relaxed) >= 1);
        tokio::task::yield_now().await;
        assert!(observed_enters.load(std::sync::atomic::Ordering::Relaxed) >= 2);
        tokio::task::yield_now().await;
        assert!(observed_enters.load(std::sync::atomic::Ordering::Relaxed) >= 3);
        return 42_u32;
    });
    let value = future.await;
    assert_eq!(value, 42_u32);
    let enter_count = enters.load(std::sync::atomic::Ordering::Relaxed);
    let exit_count = exits.load(std::sync::atomic::Ordering::Relaxed);
    assert!(enter_count >= 3);
    assert_eq!(enter_count, exit_count);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn instrumented_spans_are_usable_on_tokio_multithread_runtime() {
    let enters = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let exits = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let first_span = test_span(std::sync::Arc::clone(&enters), std::sync::Arc::clone(&exits));
    let second_span = test_span(std::sync::Arc::clone(&enters), std::sync::Arc::clone(&exits));
    let first_task = tokio::spawn(ksp_logging_lib::instrument(first_span, async {
        for _iteration in 0..32 {
            tokio::task::yield_now().await;
        }
        return 20_u32;
    }));
    let second_task = tokio::spawn(ksp_logging_lib::instrument(second_span, async {
        for _iteration in 0..32 {
            tokio::task::yield_now().await;
        }
        return 22_u32;
    }));
    let first_result = first_task.await;
    assert!(first_result.is_ok(), "first Tokio task must complete successfully");
    let first_value = match first_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(_) => return,
    };
    let second_result = second_task.await;
    assert!(second_result.is_ok(), "second Tokio task must complete successfully");
    let second_value = match second_result {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(_) => return,
    };
    assert_eq!(first_value + second_value, 42_u32);
    let enter_count = enters.load(std::sync::atomic::Ordering::Relaxed);
    let exit_count = exits.load(std::sync::atomic::Ordering::Relaxed);
    assert!(enter_count >= 4);
    assert_eq!(enter_count, exit_count);
}