use executor::praetor::{Scope, ThreadPool};

use executor::util::tree::Tree;

const TREE_SIZE: usize = 16;

fn join_scope(tree_size: usize) {
    let pool = ThreadPool::new();

    let tree = Tree::new(tree_size, 1);

    fn sum(tree: &Tree<u32>, node: usize, scope: &Scope) -> u32 {
        let node = tree.get(node);
        let (l, r) = scope.join(
            |s| node.left.map(|node| sum(tree, node, s)).unwrap_or_default(),
            |s| {
                node.right
                    .map(|node| sum(tree, node, s))
                    .unwrap_or_default()
            },
        );

        // eprintln!("node: {node:?}, l: {l}, r: {r}");

        node.leaf + l + r
    }

    for _ in 0..1000 {
        let sum = pool.scope(|s| sum(&tree, tree.root().unwrap(), s));
        std::hint::black_box(sum);
    }
}

fn join_pool(tree_size: usize) {
    let pool = ThreadPool::new();

    let tree = Tree::new(tree_size, 1);

    fn sum(tree: &Tree<u32>, node: usize, pool: &ThreadPool) -> u32 {
        let node = tree.get(node);
        let (l, r) = pool.join(
            || {
                node.left
                    .map(|node| sum(tree, node, pool))
                    .unwrap_or_default()
            },
            || {
                node.right
                    .map(|node| sum(tree, node, pool))
                    .unwrap_or_default()
            },
        );

        // eprintln!("node: {node:?}, l: {l}, r: {r}");

        node.leaf + l + r
    }

    let sum = sum(&tree, tree.root().unwrap(), &pool);
    eprintln!("sum: {sum}");
}

fn join_distaff(tree_size: usize) {
    use distaff::*;
    let pool = ThreadPool::new_with_threads(6);

    let tree = Tree::new(tree_size, 1);

    fn sum<'scope, 'env>(tree: &Tree<u32>, node: usize, scope: Scope<'scope, 'env>) -> u32 {
        let node = tree.get(node);
        let (l, r) = scope.join(
            |s| node.left.map(|node| sum(tree, node, s)).unwrap_or_default(),
            |s| {
                node.right
                    .map(|node| sum(tree, node, s))
                    .unwrap_or_default()
            },
        );

        // eprintln!("node: {node:?}, l: {l}, r: {r}");

        node.leaf + l + r
    }

    for _ in 0..1000 {
        let sum = pool.scope(|s| {
            let sum = sum(&tree, tree.root().unwrap(), s);
            sum
        });
        std::hint::black_box(sum);
    }
}

fn join_chili(tree_size: usize) {
    let tree = Tree::new(tree_size, 1u32);

    fn sum(tree: &Tree<u32>, node: usize, scope: &mut chili::Scope<'_>) -> u32 {
        let node = tree.get(node);
        let (l, r) = scope.join(
            |s| node.left.map(|node| sum(tree, node, s)).unwrap_or_default(),
            |s| {
                node.right
                    .map(|node| sum(tree, node, s))
                    .unwrap_or_default()
            },
        );

        node.leaf + l + r
    }

    for _ in 0..1000 {
        sum(&tree, tree.root().unwrap(), &mut chili::Scope::global());
        std::hint::black_box(sum);
    }
}

fn join_rayon(tree_size: usize) {
    let tree = Tree::new(tree_size, 1u32);

    fn sum(tree: &Tree<u32>, node: usize) -> u32 {
        let node = tree.get(node);
        let (l, r) = rayon::join(
            || node.left.map(|node| sum(tree, node)).unwrap_or_default(),
            || node.right.map(|node| sum(tree, node)).unwrap_or_default(),
        );

        node.leaf + l + r
    }

    for _ in 0..1000 {
        let sum = sum(&tree, tree.root().unwrap());
        std::hint::black_box(sum);
    }
}

fn main() {
    //tracing_subscriber::fmt::init();
    use tracing_subscriber::layer::SubscriberExt;
    tracing::subscriber::set_global_default(
        tracing_subscriber::registry().with(tracing_tracy::TracyLayer::default()),
    )
    .expect("Failed to set global default subscriber");

    eprintln!("Press Enter to start profiling...");
    std::io::stdin().read_line(&mut String::new()).unwrap();

    let size = std::env::args()
        .nth(2)
        .and_then(|s| s.parse::<usize>().ok())
        .unwrap_or(TREE_SIZE);

    match std::env::args().nth(1).as_deref() {
        Some("scope") => join_scope(size),
        Some("pool") => join_pool(size),
        Some("chili") => join_chili(size),
        Some("distaff") => join_distaff(size),
        Some("rayon") => join_rayon(size),
        _ => {
            eprintln!(
                "Usage: {} [scope|pool|chili|distaff|rayon] <tree_size={}>",
                std::env::args().next().unwrap(),
                TREE_SIZE
            );
            return;
        }
    }

    eprintln!("Done!");
    // // wait for user input before exiting
    // std::io::stdin().read_line(&mut String::new()).unwrap();
}