use crate::protocol::*;
use crate::protocol::eval::*;
use crate::runtime2::runtime::*;
use crate::runtime2::component::{CompCtx, CompPDL};

fn create_component(rt: &Runtime, module_name: &str, routine_name: &str, args: ValueGroup) {
    let prompt = rt.inner.protocol.new_component(
        module_name.as_bytes(), routine_name.as_bytes(), args
    ).expect("create prompt");
    let reserved = rt.inner.start_create_pdl_component();
    let ctx = CompCtx::new(&reserved);
    let component = Box::new(CompPDL::new(prompt, 0));
    let (key, _) = rt.inner.finish_create_pdl_component(reserved, component, ctx, false);
    rt.inner.enqueue_work(key);
}

fn no_args() -> ValueGroup { ValueGroup::new_stack(Vec::new()) }

#[test]
fn test_component_creation() {
    let pd = ProtocolDescription::parse(b"
    primitive nothing_at_all() {
        s32 a = 5;
        auto b = 5 + a;
    }
    ").expect("compilation");
    let rt = Runtime::new(1, true, pd);

    for _i in 0..20 {
        create_component(&rt, "", "nothing_at_all", no_args());
    }
}

#[test]
fn test_component_communication() {
    let pd = ProtocolDescription::parse(b"
    primitive sender(out<u32> o, u32 outside_loops, u32 inside_loops) {
        u32 outside_index = 0;
        while (outside_index < outside_loops) {
            u32 inside_index = 0;
            sync while (inside_index < inside_loops) {
                put(o, inside_index);
                inside_index += 1;
            }
            outside_index += 1;
        }
    }

    primitive receiver(in<u32> i, u32 outside_loops, u32 inside_loops) {
        u32 outside_index = 0;
        while (outside_index < outside_loops) {
            u32 inside_index = 0;
            sync while (inside_index < inside_loops) {
                auto val = get(i);
                while (val != inside_index) {} // infinite loop if incorrect value is received
                inside_index += 1;
            }
            outside_index += 1;
        }
    }

    composite constructor() {
        channel o_orom -> i_orom;
        channel o_mrom -> i_mrom;
        channel o_ormm -> i_ormm;
        channel o_mrmm -> i_mrmm;

        // one round, one message per round
        new sender(o_orom, 1, 1);
        new receiver(i_orom, 1, 1);

        // multiple rounds, one message per round
        new sender(o_mrom, 5, 1);
        new receiver(i_mrom, 5, 1);

        // one round, multiple messages per round
        new sender(o_ormm, 1, 5);
        new receiver(i_ormm, 1, 5);

        // multiple rounds, multiple messages per round
        new sender(o_mrmm, 5, 5);
        new receiver(i_mrmm, 5, 5);
    }").expect("compilation");
    let rt = Runtime::new(3, true, pd);
    create_component(&rt, "", "constructor", no_args());
}

#[test]
fn test_intermediate_messenger() {
    let pd = ProtocolDescription::parse(b"
    primitive receiver<T>(in<T> rx, u32 num) {
        auto index = 0;
        while (index < num) {
            sync { auto v = get(rx); }
            index += 1;
        }
    }

    primitive middleman<T>(in<T> rx, out<T> tx, u32 num) {
        auto index = 0;
        while (index < num) {
            sync { put(tx, get(rx)); }
            index += 1;
        }
    }

    primitive sender<T>(out<T> tx, u32 num) {
        auto index = 0;
        while (index < num) {
            sync put(tx, 1337);
            index += 1;
        }
    }

    composite constructor_template<T>() {
        auto num = 0;
        channel<T> tx_a -> rx_a;
        channel tx_b -> rx_b;
        new sender(tx_a, 3);
        new middleman(rx_a, tx_b, 3);
        new receiver(rx_b, 3);
    }

    composite constructor() {
        new constructor_template<u16>();
        new constructor_template<u32>();
        new constructor_template<u64>();
        new constructor_template<s16>();
        new constructor_template<s32>();
        new constructor_template<s64>();
    }
    ").expect("compilation");
    let rt = Runtime::new(3, true, pd);
    create_component(&rt, "", "constructor", no_args());
}

#[test]
fn test_simple_select() {
    let pd = ProtocolDescription::parse(b"
    func infinite_assert<T>(T val, T expected) -> () {
        while (val != expected) { print(\"nope!\"); }
        return ();
    }

    primitive receiver(in<u32> in_a, in<u32> in_b, u32 num_sends) {
        auto num_from_a = 0;
        auto num_from_b = 0;
        while (num_from_a + num_from_b < 2 * num_sends) {
            sync select {
                auto v = get(in_a) -> {
                    print(\"got something from A\");
                    auto _ = infinite_assert(v, num_from_a);
                    num_from_a += 1;
                }
                auto v = get(in_b) -> {
                    print(\"got something from B\");
                    auto _ = infinite_assert(v, num_from_b);
                    num_from_b += 1;
                }
            }
        }
    }

    primitive sender(out<u32> tx, u32 num_sends) {
        auto index = 0;
        while (index < num_sends) {
            sync {
                put(tx, index);
                index += 1;
            }
        }
    }

    composite constructor() {
        auto num_sends = 15;
        channel tx_a -> rx_a;
        channel tx_b -> rx_b;
        new sender(tx_a, num_sends);
        new receiver(rx_a, rx_b, num_sends);
        new sender(tx_b, num_sends);
    }
    ").expect("compilation");
    let rt = Runtime::new(3, false, pd);
    create_component(&rt, "", "constructor", no_args());
}

#[test]
fn test_unguarded_select() {
    let pd = ProtocolDescription::parse(b"
    primitive constructor_outside_select() {
        u32 index = 0;
        while (index < 5) {
            sync select { auto v = () -> print(\"hello\"); }
            index += 1;
        }
    }

    primitive constructor_inside_select() {
        u32 index = 0;
        while (index < 5) {
            sync select { auto v = () -> index += 1; }
        }
    }
    ").expect("compilation");
    let rt = Runtime::new(3, false, pd);
    create_component(&rt, "", "constructor_outside_select", no_args());
    create_component(&rt, "", "constructor_inside_select", no_args());
}

#[test]
fn test_empty_select() {
    let pd = ProtocolDescription::parse(b"
    primitive constructor() {
        u32 index = 0;
        while (index < 5) {
            sync select {}
            index += 1;
        }
    }
    ").expect("compilation");
    let rt = Runtime::new(3, false, pd);
    create_component(&rt, "", "constructor", no_args());
}

#[test]
fn test_random_u32_temporary_thingo() {
    let pd = ProtocolDescription::parse(b"
    primitive random_taker(in<u32> generator) {
        sync {
            auto a = get(generator);
        }
    }

    composite constructor() {
        channel tx -> rx;
        new random_u32(tx, 1, 100);
        new random_taker(rx);
    }
    ").expect("compilation");
    let rt = Runtime::new(1, true, pd);
    create_component(&rt, "", "constructor", no_args());
}