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

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

const TREE_SIZE: usize = 16;

fn join_scope() {
    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() {
    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() {
    use distaff::*;
    let pool = ThreadPool::new();
    let tree = Tree::new(TREE_SIZE, 1);

    fn sum<'scope, 'env>(tree: &Tree<u32>, node: usize, scope: &'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
        });
        eprintln!("sum: {sum}");
        std::hint::black_box(sum);
    }
}

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