//! Concurrent access tests

use dagx::{task, DagRunner};
use dagx_test::task_fn;

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::task::JoinSet;

struct SleepTask(usize);

#[task]
impl SleepTask {
    async fn run(&self) -> usize {
        tokio::time::sleep(Duration::from_millis(5)).await;
        self.0
    }
}

#[tokio::test]
async fn test_concurrent_run_calls() {
    // Test that multiple run() calls are handled correctly
    let dag = Arc::new(DagRunner::new());

    // Add some tasks
    let tasks: Vec<_> = (0..10).map(|i| dag.add_task(SleepTask(i))).collect();

    // Try to run the DAG multiple times concurrently
    let mut handles = JoinSet::new();

    for _ in 0..5 {
        let dag = Arc::clone(&dag);
        handles.spawn(async move {
            dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
                .await
        });
    }

    // Only one run should succeed, others should fail with concurrent execution error
    let mut success_count = 0;
    let mut error_count = 0;

    while let Some(result) = handles.join_next().await {
        match result.unwrap() {
            Ok(_) => success_count += 1,
            Err(_) => error_count += 1,
        }
    }

    // Exactly one should succeed, the rest should fail
    assert_eq!(success_count, 1);
    assert_eq!(error_count, 4);

    // Verify results are consistent from the successful run
    for (i, task) in tasks.into_iter().enumerate() {
        assert_eq!(dag.get(task).unwrap(), i);
    }
}

#[tokio::test]
async fn test_concurrent_execution_and_building() {
    // Test that we can build and execute concurrently
    let dag = Arc::new(DagRunner::new());

    // Start with some initial tasks
    let initial: Vec<_> = (0..10)
        .map(|i| dag.add_task(task_fn::<(), _, _>(move |_: ()| i)))
        .collect();

    // Execute the initial DAG
    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await
        .unwrap();

    // Now try to add more tasks while also reading results
    let mut handles = JoinSet::new();

    // Reader tasks
    for task in initial {
        let dag = Arc::clone(&dag);
        let task_handle: dagx::TaskHandle<i32> = task.into();
        handles.spawn(async move { dag.get(task_handle).unwrap() });
    }

    // Builder tasks (adding new nodes)
    for i in 10..20 {
        let dag = Arc::clone(&dag);
        handles.spawn(async move {
            dag.add_task(task_fn::<(), _, _>(move |_: ()| i));
            i
        });
    }

    // Wait for all operations
    while let Some(result) = handles.join_next().await {
        result.unwrap();
    }

    // Execute again with the new tasks
    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await
        .unwrap();
}

#[tokio::test]
async fn test_dag_builder_thread_safety() {
    // Test that the DAG builder is thread-safe
    let dag = Arc::new(DagRunner::new());
    let counter = Arc::new(AtomicUsize::new(0));

    let mut handles = JoinSet::new();

    // Spawn tasks that build complex sub-graphs concurrently
    for thread_id in 0..10 {
        let dag = Arc::clone(&dag);
        let counter = Arc::clone(&counter);

        handles.spawn(async move {
            // Each thread creates a small pipeline
            let source = dag.add_task(task_fn::<(), _, _>(move |_: ()| thread_id));

            let transform1 = dag
                .add_task(task_fn::<i32, _, _>(move |&x: &i32| x * 10))
                .depends_on(source);

            let transform2 = dag
                .add_task(task_fn::<i32, _, _>(move |&x: &i32| x + thread_id))
                .depends_on(transform1);

            counter.fetch_add(1, Ordering::SeqCst);
            transform2
        });
    }

    let mut final_tasks = vec![];
    while let Some(result) = handles.join_next().await {
        final_tasks.push(result.unwrap());
    }

    // Verify all threads completed
    assert_eq!(counter.load(Ordering::SeqCst), 10);

    // Execute the combined DAG
    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await
        .unwrap();

    // Verify each pipeline computed correctly
    for task in final_tasks.iter() {
        let result = dag.get(task).unwrap();
        // Each pipeline: thread_id * 10 + thread_id
        assert!((0..100).contains(&result));
    }
}

#[tokio::test]
async fn test_concurrent_get_operations() {
    // Test concurrent get() operations
    let dag = Arc::new(DagRunner::new());

    // Create and execute tasks
    let tasks: Vec<_> = (0..100)
        .map(|i| dag.add_task(task_fn::<(), _, _>(move |_: ()| i)))
        .collect();

    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await
        .unwrap();

    // Now access results concurrently
    let mut handles = JoinSet::new();

    for task in tasks.into_iter() {
        let dag = Arc::clone(&dag);
        let task_handle: dagx::TaskHandle<i32> = task.into();
        handles.spawn(async move { dag.get(task_handle).unwrap() });
    }

    // Collect and verify results
    let mut results = vec![];
    while let Some(result) = handles.join_next().await {
        results.push(result.unwrap());
    }

    // Results should contain all values 0..100
    results.sort();
    for (i, val) in results.iter().enumerate() {
        assert_eq!(*val, i as i32);
    }
}

#[tokio::test]
async fn test_no_race_in_dependency_tracking() {
    // Test that concurrent dependency modifications don't cause races
    let dag = Arc::new(DagRunner::new());
    let source = dag.add_task(task_fn::<(), _, _>(|_: ()| 42));

    let mut handles = JoinSet::new();

    // Multiple threads create tasks depending on the same source
    let source_handle: dagx::TaskHandle<i32> = source.into();
    for i in 0..50 {
        let dag = Arc::clone(&dag);
        let source_clone = source_handle;

        handles.spawn(async move {
            dag.add_task(task_fn::<i32, _, _>(move |&x: &i32| x + i))
                .depends_on(source_clone)
        });
    }

    let mut dependents = vec![];
    while let Some(result) = handles.join_next().await {
        dependents.push(result.unwrap());
    }

    // Execute
    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await
        .unwrap();

    // All dependents should have correct values
    for dep in dependents.iter() {
        let result = dag.get(dep).unwrap();
        assert!((42..92).contains(&result)); // 42 + (0..50)
    }
}

#[tokio::test]
async fn test_high_concurrency_stress() {
    // Stress test with high concurrency
    let dag = Arc::new(DagRunner::new());
    let completed = Arc::new(AtomicUsize::new(0));

    let mut handles = JoinSet::new();

    // 1000 concurrent operations
    for i in 0..1000 {
        let dag = Arc::clone(&dag);
        let completed = Arc::clone(&completed);

        handles.spawn(async move {
            // Mix of operations
            match i % 3 {
                0 => {
                    // Add a source task
                    dag.add_task(task_fn::<(), _, _>(move |_: ()| i));
                }
                1 => {
                    // Add a task with a simple dependency
                    let source = dag.add_task(task_fn::<(), _, _>(move |_: ()| i));
                    dag.add_task(task_fn::<i32, _, _>(move |&x: &i32| x * 2))
                        .depends_on(source);
                }
                2 => {
                    // Add a more complex subgraph
                    let a = dag.add_task(task_fn::<(), _, _>(move |_: ()| i));
                    let b = dag.add_task(task_fn::<(), _, _>(move |_: ()| i + 1));
                    dag.add_task(task_fn::<(i32, i32), _, _>(move |(x, y): (&i32, &i32)| {
                        x + y
                    }))
                    .depends_on((a, b));
                }
                _ => unreachable!(),
            }
            completed.fetch_add(1, Ordering::SeqCst);
        });
    }

    // Wait for all operations to complete
    while let Some(result) = handles.join_next().await {
        result.unwrap();
    }

    assert_eq!(completed.load(Ordering::SeqCst), 1000);

    // The DAG should still be executable
    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await
        .unwrap();
}

#[tokio::test]
async fn test_lock_contention_measurement() {
    // Test to measure lock contention under concurrent load
    let dag = Arc::new(DagRunner::new());
    let start = tokio::time::Instant::now();
    let operations = Arc::new(AtomicUsize::new(0));

    let mut handles = JoinSet::new();

    // Create many concurrent operations that will contend for locks
    for i in 0..100 {
        let dag = Arc::clone(&dag);
        let ops = Arc::clone(&operations);

        handles.spawn(async move {
            for j in 0..10 {
                dag.add_task(task_fn::<(), _, _>(move |_: ()| i * 10 + j));
                ops.fetch_add(1, Ordering::SeqCst);

                // Small yield to increase contention likelihood
                tokio::task::yield_now().await;
            }
        });
    }

    // Wait for completion
    while let Some(result) = handles.join_next().await {
        result.unwrap();
    }

    let duration = start.elapsed();
    let total_ops = operations.load(Ordering::SeqCst);

    println!(
        "Completed {} operations in {:?} ({:.0} ops/sec)",
        total_ops,
        duration,
        total_ops as f64 / duration.as_secs_f64()
    );

    // Should complete reasonably quickly even with contention
    assert!(
        duration < Duration::from_secs(5),
        "Too much lock contention"
    );
}

#[tokio::test]
async fn test_atomic_dag_operations() {
    // Test that DAG operations are atomic (all-or-nothing)
    let dag = Arc::new(DagRunner::new());

    // Create a base task
    let base = dag.add_task(task_fn::<(), _, _>(|_: ()| 100));

    // Try to create dependent tasks concurrently
    let mut handles = JoinSet::new();
    let success_count = Arc::new(AtomicUsize::new(0));

    let base_handle: dagx::TaskHandle<i32> = base.into();
    for i in 0..50 {
        let dag = Arc::clone(&dag);
        let base_clone = base_handle;
        let success = Arc::clone(&success_count);

        handles.spawn(async move {
            let dep = dag
                .add_task(task_fn::<i32, _, _>(move |&x: &i32| x + i))
                .depends_on(base_clone);

            success.fetch_add(1, Ordering::SeqCst);
            dep
        });
    }

    let mut deps = vec![];
    while let Some(result) = handles.join_next().await {
        deps.push(result.unwrap());
    }

    // All operations should have succeeded
    assert_eq!(success_count.load(Ordering::SeqCst), 50);
    assert_eq!(deps.len(), 50);

    // Execute and verify
    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await
        .unwrap();

    for dep in deps.iter() {
        let result = dag.get(dep).unwrap();
        assert!((100..150).contains(&result));
    }
}

#[tokio::test]
async fn test_no_deadlock_on_circular_waiting() {
    // Test that we don't deadlock even with complex concurrent operations
    let dag = Arc::new(DagRunner::new());

    // Create multiple threads that might cause circular waiting patterns
    let mut handles = JoinSet::new();

    for i in 0..10 {
        let dag = Arc::clone(&dag);

        handles.spawn(async move {
            // Each thread creates a chain of tasks
            let first = dag.add_task(task_fn::<(), _, _>(move |_: ()| i));
            let mut chain = vec![first.into()];

            for j in 0..5 {
                let next = dag
                    .add_task(task_fn::<i32, _, _>(move |&x: &i32| x + j))
                    .depends_on(chain.last().unwrap());
                chain.push(next);

                // Yield to potentially cause interleaving
                tokio::task::yield_now().await;
            }

            *chain.last().unwrap()
        });
    }

    // Set a timeout to detect deadlocks
    let result = tokio::time::timeout(Duration::from_secs(5), async {
        let mut finals = vec![];
        while let Some(result) = handles.join_next().await {
            finals.push(result.unwrap());
        }
        finals
    })
    .await;

    assert!(result.is_ok(), "Deadlock detected - operation timed out");

    // Should be able to execute the DAG
    dag.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
        .await
        .unwrap();
}