use super::*;
use crate::{PortId, ProtocolDescription};
use crate::common::Id;
use crate::protocol::eval::*;

const NUM_THREADS: u32 = 3;     // number of threads in runtime
const NUM_INSTANCES: u32 = 5;   // number of test instances constructed
const NUM_LOOPS: u32 = 5;       // number of loops within a single test (not used by all tests)

fn create_runtime(pdl: &str) -> Runtime {
    let protocol = ProtocolDescription::parse(pdl.as_bytes()).expect("parse pdl");
    let runtime = Runtime::new(NUM_THREADS, protocol);

    return runtime;
}

fn run_test_in_runtime<F: Fn(&mut ApplicationInterface)>(pdl: &str, constructor: F) {
    let protocol = ProtocolDescription::parse(pdl.as_bytes())
        .expect("parse PDL");
    let runtime = Runtime::new(NUM_THREADS, protocol);

    let mut api = runtime.create_interface();
    for _ in 0..NUM_INSTANCES {
        constructor(&mut api);
    }

    // Wait until done :)
}

struct TestTimer {
    name: &'static str,
    started: std::time::Instant
}

impl TestTimer {
    fn new(name: &'static str) -> Self {
        Self{ name, started: std::time::Instant::now() }
    }
}

impl Drop for TestTimer {
    fn drop(&mut self) {
        let delta = std::time::Instant::now() - self.started;
        let nanos = (delta.as_secs_f64() * 1_000_000.0) as u64;
        let millis = nanos / 1000;
        let nanos = nanos % 1000;
        println!("[{}] Took {:>4}.{:03} ms", self.name, millis, nanos);
    }
}

#[test]
fn test_put_and_get() {
    const CODE: &'static str = "
    primitive putter(out<bool> sender, u32 loops) {
        u32 index = 0;
        while (index < loops) {
            synchronous {
                put(sender, true);
            }
            index += 1;
        }
    }

    primitive getter(in<bool> receiver, u32 loops) {
        u32 index = 0;
        while (index < loops) {
            synchronous {
                auto result = get(receiver);
                assert(result);
            }
            index += 1;
        }
    }
    ";

    let thing = TestTimer::new("put_and_get");
    run_test_in_runtime(CODE, |api| {
        let channel = api.create_channel();

        api.create_connector("", "putter", ValueGroup::new_stack(vec![
            Value::Output(PortId(Id{ connector_id: 0, u32_suffix: channel.putter_id.index })),
            Value::UInt32(NUM_LOOPS)
        ])).expect("create putter");

        api.create_connector("", "getter", ValueGroup::new_stack(vec![
            Value::Input(PortId(Id{ connector_id: 0, u32_suffix: channel.getter_id.index })),
            Value::UInt32(NUM_LOOPS)
        ])).expect("create getter");
    });
}

#[test]
fn test_star_shaped_request() {
    const CODE: &'static str = "
    primitive edge(in<u32> input, out<u32> output, u32 loops) {
        u32 index = 0;
        while (index < loops) {
            synchronous {
                auto req = get(input);
                put(output, req * 2);
            }
            index += 1;
        }
    }

    primitive center(out<u32>[] requests, in<u32>[] responses, u32 loops) {
        u32 loop_index = 0;
        auto num_edges = length(requests);

        while (loop_index < loops) {
            // print(\"starting loop\");
            synchronous {
                u32 edge_index = 0;
                u32 sum = 0;
                while (edge_index < num_edges) {
                    put(requests[edge_index], edge_index);
                    auto response = get(responses[edge_index]);
                    sum += response;
                    edge_index += 1;
                }

                assert(sum == num_edges * (num_edges - 1));
            }
            // print(\"ending loop\");
            loop_index += 1;
        }
    }

    composite constructor(u32 num_edges, u32 num_loops) {
        auto requests = {};
        auto responses = {};

        u32 edge_index = 0;
        while (edge_index < num_edges) {
            channel req_put -> req_get;
            channel resp_put -> resp_get;
            new edge(req_get, resp_put, num_loops);
            requests @= { req_put };
            responses @= { resp_get };

            edge_index += 1;
        }

        new center(requests, responses, num_loops);
    }
    ";

    let thing = TestTimer::new("star_shaped_request");
    run_test_in_runtime(CODE, |api| {
        api.create_connector("", "constructor", ValueGroup::new_stack(vec![
            Value::UInt32(5),
            Value::UInt32(NUM_LOOPS),
        ]));
    });
}

#[test]
fn test_conga_line_request() {
    const CODE: &'static str = "
    primitive start(out<u32> req, in<u32> resp, u32 num_nodes, u32 num_loops) {
        u32 loop_index = 0;
        u32 initial_value = 1337;
        while (loop_index < num_loops) {
            synchronous {
                put(req, initial_value);
                auto result = get(resp);
                assert(result == initial_value + num_nodes * 2);
            }
            loop_index += 1;
        }
    }

    primitive middle(
        in<u32> req_in, out<u32> req_forward,
        in<u32> resp_in, out<u32> resp_forward,
        u32 num_loops
    ) {
        u32 loop_index = 0;
        while (loop_index < num_loops) {
            synchronous {
                auto req = get(req_in);
                put(req_forward, req + 1);
                auto resp = get(resp_in);
                put(resp_forward, resp + 1);
            }
            loop_index += 1;
        }
    }

    primitive end(in<u32> req_in, out<u32> resp_out, u32 num_loops) {
        u32 loop_index = 0;
        while (loop_index < num_loops) {
            synchronous {
                auto req = get(req_in);
                put(resp_out, req);
            }
            loop_index += 1;
        }
    }

    composite constructor(u32 num_nodes, u32 num_loops) {
        channel initial_req -> req_in;
        channel resp_out -> final_resp;
        new start(initial_req, final_resp, num_nodes, num_loops);

        in<u32> last_req_in = req_in;
        out<u32> last_resp_out = resp_out;

        u32 node = 0;
        while (node < num_nodes) {
            channel new_req_fw -> new_req_in;
            channel new_resp_out -> new_resp_in;
            new middle(last_req_in, new_req_fw, new_resp_in, last_resp_out, num_loops);

            last_req_in = new_req_in;
            last_resp_out = new_resp_out;

            node += 1;
        }

        new end(last_req_in, last_resp_out, num_loops);
    }
    ";

    let thing = TestTimer::new("conga_line_request");
    run_test_in_runtime(CODE, |api| {
        api.create_connector("", "constructor", ValueGroup::new_stack(vec![
            Value::UInt32(5),
            Value::UInt32(NUM_LOOPS)
        ]));
    });
}