//! Complex DAG pattern tests

use crate::common::task_fn;
use dagx::{DagResult, DagRunner, TaskHandle};

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

#[tokio::test]
async fn test_workflow_with_conditional_paths() -> DagResult<()> {
    let dag = DagRunner::new();

    // Input validation
    let input = dag.add_task(task_fn::<(), _, _>(|_: ()| 42));

    let validate = dag
        .add_task(task_fn::<i32, _, _>(|&x: &i32| {
            if x > 0 && x < 100 {
                Ok(x)
            } else {
                Err("Invalid input")
            }
        }))
        .depends_on(input);

    // Process valid input
    let process = dag
        .add_task(task_fn::<Result<_, _>, _, _>(
            |result: &Result<i32, &str>| match result {
                Ok(x) => *x * 2,
                Err(_) => -1,
            },
        ))
        .depends_on(validate);

    // Parallel enhancement steps
    let enhance1 = dag
        .add_task(task_fn::<i32, _, _>(|&x: &i32| x + 10))
        .depends_on(process);

    let enhance2 = dag
        .add_task(task_fn::<i32, _, _>(|&x: &i32| x * 3))
        .depends_on(process);

    let enhance3 = dag
        .add_task(task_fn::<i32, _, _>(|&x: &i32| x - 5))
        .depends_on(process);

    // Combine enhancements
    let combine = dag
        .add_task(task_fn::<(i32, i32, i32), _, _>(
            |(&e1, &e2, &e3): (&i32, &i32, &i32)| vec![e1, e2, e3],
        ))
        .depends_on((&enhance1, &enhance2, &enhance3));

    // Final aggregation
    let final_result = dag
        .add_task(task_fn::<Vec<_>, _, _>(|values: &Vec<i32>| {
            values.iter().sum::<i32>()
        }))
        .depends_on(combine);

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

    assert_eq!(dag.get(process)?, 84);
    let result = dag.get(final_result)?;
    assert!(result > 0);

    Ok(())
}

#[tokio::test]
async fn test_map_reduce_pattern() -> DagResult<()> {
    let dag = DagRunner::new();

    // Map phase: process 16 inputs in parallel
    let mappers: Vec<dagx::TaskHandle<Vec<(&str, usize)>>> = (0..16)
        .map(|i| {
            let task = dag.add_task(task_fn::<(), _, _>(move |_: ()| {
                // Simulate word count
                vec![("word1", i % 3), ("word2", i % 5), ("word3", i % 7)]
            }));
            task.into()
        })
        .collect();

    // Shuffle phase: group by key (simplified - just 4 reducers)
    let shuffled: Vec<_> = (0..4)
        .map(|reducer_id| {
            // Each reducer gets data from 4 mappers
            let mapper_indices: Vec<usize> = (0..4).map(|i| reducer_id * 4 + i).collect();

            let deps: Vec<_> = mapper_indices.iter().map(|&i| mappers[i]).collect();

            match deps.len() {
                4 => dag
                    .add_task(task_fn::<
                        (
                            Vec<(&str, usize)>,
                            Vec<(&str, usize)>,
                            Vec<(&str, usize)>,
                            Vec<(&str, usize)>,
                        ),
                        _,
                        _,
                    >(
                        move |inputs: (
                            &Vec<(&str, usize)>,
                            &Vec<(&str, usize)>,
                            &Vec<(&str, usize)>,
                            &Vec<(&str, usize)>,
                        )| {
                            let mut combined = Vec::<(&str, usize)>::new();
                            combined.extend(inputs.0);
                            combined.extend(inputs.1);
                            combined.extend(inputs.2);
                            combined.extend(inputs.3);
                            combined
                        },
                    ))
                    .depends_on((&deps[0], &deps[1], &deps[2], &deps[3])),
                _ => unreachable!(),
            }
        })
        .collect();

    // Reduce phase
    let reducers: Vec<_> = shuffled
        .iter()
        .map(|shuffled_data| {
            dag.add_task(task_fn::<Vec<_>, _, _>(|data: &Vec<(&str, usize)>| {
                // Sum values by key
                let mut totals = std::collections::HashMap::new();
                for (key, value) in data {
                    *totals.entry(*key).or_insert(0) += value;
                }
                totals.len()
            }))
            .depends_on(shuffled_data)
        })
        .collect();

    // Final aggregation
    let final_task = dag
        .add_task(task_fn::<(_, _, _, _), _, _>(
            |(r1, r2, r3, r4): (&usize, &usize, &usize, &usize)| r1 + r2 + r3 + r4,
        ))
        .depends_on((&reducers[0], &reducers[1], &reducers[2], &reducers[3]));

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

    let result = dag.get(final_task)?;
    assert!(result > 0);

    Ok(())
}

#[tokio::test]
async fn test_pipeline_with_feedback() -> DagResult<()> {
    let dag = DagRunner::new();
    let iteration_count = Arc::new(AtomicUsize::new(0));

    // Initial value
    let init = dag.add_task(task_fn::<(), _, _>(|_: ()| 1));

    // Stage 1: Transform
    let stage1 = dag
        .add_task(task_fn::<i32, _, _>({
            let counter = iteration_count.clone();
            move |x: &i32| {
                let counter = counter.clone();
                counter.fetch_add(1, Ordering::SeqCst);
                x * 2
            }
        }))
        .depends_on(init);

    // Stage 2: Filter
    let stage2 = dag
        .add_task(task_fn::<i32, _, _>(
            |&x: &i32| if x < 100 { Some(x) } else { None },
        ))
        .depends_on(stage1);

    // Stage 3: Accumulate
    let stage3 = dag
        .add_task(task_fn::<Option<_>, _, _>(|opt: &Option<i32>| {
            opt.map(|x| x + 10).unwrap_or(0)
        }))
        .depends_on(stage2);

    // Note: Can't create actual cycles in DAG, but we simulate feedback
    // by having a separate "feedback" task that would feed back in a real system
    let _feedback = dag
        .add_task(task_fn::<i32, _, _>(|&x: &i32| {
            // In a real feedback loop, this would feed back to stage1
            x / 2
        }))
        .depends_on(stage3);

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

    assert_eq!(dag.get(stage1)?, 2);
    assert_eq!(dag.get(stage3)?, 12);
    assert_eq!(iteration_count.load(Ordering::SeqCst), 1);

    Ok(())
}

#[tokio::test]
async fn test_scatter_gather_pattern() -> DagResult<()> {
    let dag = DagRunner::new();

    // Scatter: distribute work to multiple workers
    let source: TaskHandle<_> = dag
        .add_task(task_fn::<(), _, _>(|_: ()| {
            vec![10, 20, 30, 40, 50, 60, 70, 80]
        }))
        .into();

    // Workers process chunks
    let workers: Vec<_> = (0..4)
        .map(|worker_id| {
            dag.add_task(task_fn::<Vec<_>, _, _>(move |data: &Vec<i32>| {
                // Each worker processes a chunk
                let chunk_size = data.len() / 4;
                let start = worker_id * chunk_size;
                let end = if worker_id == 3 {
                    data.len()
                } else {
                    start + chunk_size
                };

                data[start..end].iter().sum::<i32>()
            }))
            .depends_on(source)
        })
        .collect();

    // Gather: collect results from all workers
    let gather = dag
        .add_task(task_fn::<(i32, i32, i32, i32), _, _>(
            |(w0, w1, w2, w3): (&i32, &i32, &i32, &i32)| w0 + w1 + w2 + w3,
        ))
        .depends_on((&workers[0], &workers[1], &workers[2], &workers[3]));

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

    // Sum of 10+20+30+40+50+60+70+80 = 360
    assert_eq!(dag.get(gather)?, 360);

    Ok(())
}

#[tokio::test]
async fn test_fork_join_with_barriers() -> DagResult<()> {
    let dag = DagRunner::new();
    let phase_counter = Arc::new(AtomicUsize::new(0));

    // Fork phase
    let parent = dag.add_task(task_fn::<(), _, _>({
        let counter = phase_counter.clone();
        move |_: ()| {
            let counter = counter.clone();
            counter.store(1, Ordering::SeqCst);
            100
        }
    }));

    // Multiple parallel paths
    let mut path1 = Vec::new();
    let parent_handle: dagx::TaskHandle<i32> = parent.into();
    let mut prev1 = parent_handle;
    for i in 0..3 {
        let counter = phase_counter.clone();
        let task = dag
            .add_task(task_fn::<i32, _, _>(move |&x: &i32| {
                let counter = counter.clone();
                assert_eq!(counter.load(Ordering::SeqCst), 1);
                x + i
            }))
            .depends_on(prev1);
        path1.push(task);
        prev1 = task;
    }

    let mut path2 = Vec::new();
    let mut prev2 = parent_handle;
    for i in 0..3 {
        let counter = phase_counter.clone();
        let task = dag
            .add_task(task_fn::<i32, _, _>(move |&x: &i32| {
                let counter = counter.clone();
                assert_eq!(counter.load(Ordering::SeqCst), 1);
                x * 2 + i
            }))
            .depends_on(prev2);
        path2.push(task);
        prev2 = task;
    }

    // Barrier: wait for all paths
    let barrier = dag
        .add_task(task_fn::<(i32, i32), _, _>({
            let counter = phase_counter.clone();
            move |(p1, p2): (&i32, &i32)| {
                let counter = counter.clone();
                counter.store(2, Ordering::SeqCst);
                p1 + p2
            }
        }))
        .depends_on((&path1[2], &path2[2]));

    // Join phase
    let join = dag
        .add_task(task_fn::<i32, _, _>({
            let counter = phase_counter.clone();
            move |&x: &i32| {
                let counter = counter.clone();
                assert_eq!(counter.load(Ordering::SeqCst), 2);
                x
            }
        }))
        .depends_on(barrier);

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

    let _ = dag.get(join)?;
    assert_eq!(phase_counter.load(Ordering::SeqCst), 2);

    Ok(())
}