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 (key, _) = rt.inner.finish_create_pdl_component(reserved, CompPDL::new(prompt, 0), 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_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());
}