/*
 * Copyright 2022 Fluence Labs Limited
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use air_interpreter_data::ExecutionTrace;
use air_test_utils::prelude::*;

use pretty_assertions::assert_eq;

#[test]
fn join_chat_1() {
    use std::collections::HashSet;

    let relay_1_peer_id = "relay_1_peer_id";
    let mut relay_1 = create_avm(unit_call_service(), relay_1_peer_id);

    let relay_2_peer_id = "relay_2_peer_id";
    let mut relay_2 = create_avm(unit_call_service(), relay_2_peer_id);

    let client_1_peer_id = "client_1_peer_id";
    let mut client_1 = create_avm(unit_call_service(), client_1_peer_id);

    let client_2_peer_id = "client_2_peer_id";
    let mut client_2 = create_avm(unit_call_service(), client_2_peer_id);

    let remote_peer_id = "remote_peer_id";
    let members = json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]);
    let members_call_service: CallServiceClosure =
        Box::new(move |_| -> CallServiceResult { CallServiceResult::ok(members.clone()) });
    let mut remote = create_avm(members_call_service, remote_peer_id);

    let script = format!(
        r#"
            (seq
                (call "{relay_1_peer_id}" ("identity" "") [] $void1)
                (seq
                    (call "{remote_peer_id}" ("552196ea-b9b2-4761-98d4-8e7dba77fac4" "add") [] $void2)
                    (seq
                        (call "{remote_peer_id}" ("920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9" "get_users") [] members)
                        (fold members m
                            (par
                                (seq
                                    (call m.$.[1]! ("identity" "") [] $void)
                                    (call m.$.[0]! ("fgemb3" "add") [] $void3)
                                )
                                (next m)
                            )
                        )
                    )
                )
            )
        "#
    );

    let client_1_result = checked_call_vm!(client_1, <_>::default(), &script, "", "");

    let client_1_actual_trace = trace_from_result(&client_1_result);
    let client_1_expected_trace = vec![executed_state::request_sent_by(client_1_peer_id)];

    assert_eq!(client_1_actual_trace, client_1_expected_trace);
    assert_eq!(client_1_result.next_peer_pks, vec![String::from(relay_1_peer_id)]);

    let relay_1_result = checked_call_vm!(relay_1, <_>::default(), &script, client_1_result.data, "");

    let relay_1_actual_trace = trace_from_result(&relay_1_result);
    let relay_1_expected_trace = vec![
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        executed_state::request_sent_by(relay_1_peer_id),
    ];

    assert_eq!(relay_1_actual_trace, relay_1_expected_trace);
    assert_eq!(relay_1_result.next_peer_pks, vec![String::from(remote_peer_id)]);

    let remote_result = checked_call_vm!(remote, <_>::default(), &script, relay_1_result.data, "");

    let remote_actual_trace = trace_from_result(&remote_result);
    let remote_expected_trace = vec![
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        stream!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            0,
            peer = remote_peer_id,
            service = "552196ea-b9b2-4761-98d4-8e7dba77fac4",
            function = "add"
        ),
        scalar!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            peer = remote_peer_id,
            service = "920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9",
            function = "get_users"
        ),
        executed_state::par(1, 2),
        executed_state::request_sent_by(remote_peer_id),
        executed_state::par(1, 0),
        executed_state::request_sent_by(remote_peer_id),
    ];

    let actual_remote_next_peer_pks: HashSet<_> = remote_result.next_peer_pks.iter().map(|s| s.as_str()).collect();
    let expected_next_peer_pks = maplit::hashset! {
        relay_1_peer_id,
        relay_2_peer_id,
    };

    assert_eq!(remote_actual_trace, remote_expected_trace);
    assert_eq!(actual_remote_next_peer_pks, expected_next_peer_pks);

    let relay_1_result = checked_call_vm!(relay_1, <_>::default(), &script, remote_result.data.clone(), "");

    let relay_1_actual_trace = trace_from_result(&relay_1_result);

    let relay_1_expected_trace = ExecutionTrace::from(vec![
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        stream!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            0,
            peer = remote_peer_id,
            service = "552196ea-b9b2-4761-98d4-8e7dba77fac4",
            function = "add"
        ),
        scalar!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            peer = remote_peer_id,
            service = "920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9",
            function = "get_users"
        ),
        executed_state::par(2, 2),
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        executed_state::request_sent_by(relay_1_peer_id),
        executed_state::par(1, 0),
        executed_state::request_sent_by(remote_peer_id),
    ]);

    assert_eq!(relay_1_actual_trace, relay_1_expected_trace);
    assert_eq!(relay_1_result.next_peer_pks, vec![String::from(client_1_peer_id)]);

    let client_1_result = checked_call_vm!(client_1, <_>::default(), &script, relay_1_result.data, "");

    let client_1_actual_trace = trace_from_result(&client_1_result);

    let client_1_expected_trace = ExecutionTrace::from(vec![
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        stream!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            0,
            peer = remote_peer_id,
            service = "552196ea-b9b2-4761-98d4-8e7dba77fac4",
            function = "add"
        ),
        scalar!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            peer = remote_peer_id,
            service = "920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9",
            function = "get_users"
        ),
        executed_state::par(2, 2),
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        stream!(
            "result from unit_call_service",
            0,
            peer = client_1_peer_id,
            service = "fgemb3",
            function = "add"
        ),
        executed_state::par(1, 0),
        executed_state::request_sent_by(remote_peer_id),
    ]);

    assert_eq!(client_1_actual_trace, client_1_expected_trace);
    assert!(client_1_result.next_peer_pks.is_empty());

    let relay_2_result = checked_call_vm!(relay_2, <_>::default(), &script, remote_result.data, "");

    let relay_2_actual_trace = trace_from_result(&relay_2_result);

    let relay_2_expected_trace = ExecutionTrace::from(vec![
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        stream!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            0,
            peer = remote_peer_id,
            service = "552196ea-b9b2-4761-98d4-8e7dba77fac4",
            function = "add"
        ),
        scalar!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            peer = remote_peer_id,
            service = "920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9",
            function = "get_users"
        ),
        executed_state::par(1, 3),
        executed_state::request_sent_by(remote_peer_id),
        executed_state::par(2, 0),
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_2_peer_id,
            service = "identity"
        ),
        executed_state::request_sent_by(relay_2_peer_id),
    ]);

    assert_eq!(relay_2_actual_trace, relay_2_expected_trace);
    assert_eq!(relay_2_result.next_peer_pks, vec![String::from(client_2_peer_id)]);

    let client_2_result = checked_call_vm!(client_2, <_>::default(), script, relay_2_result.data, "");

    let client_2_actual_trace = trace_from_result(&client_2_result);

    let client_2_expected_trace = ExecutionTrace::from(vec![
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        stream!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            0,
            peer = remote_peer_id,
            service = "552196ea-b9b2-4761-98d4-8e7dba77fac4",
            function = "add"
        ),
        scalar!(
            json!([[client_1_peer_id, relay_1_peer_id], [client_2_peer_id, relay_2_peer_id]]),
            peer = remote_peer_id,
            service = "920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9",
            function = "get_users"
        ),
        executed_state::par(1, 3),
        executed_state::request_sent_by(remote_peer_id),
        executed_state::par(2, 0),
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_2_peer_id,
            service = "identity"
        ),
        stream!(
            "result from unit_call_service",
            0,
            peer = client_2_peer_id,
            service = "fgemb3",
            function = "add"
        ),
    ]);

    assert_eq!(client_2_actual_trace, client_2_expected_trace);
    assert!(client_2_result.next_peer_pks.is_empty());
}

#[test]
fn join_chat_2() {
    let members_call_service1: CallServiceClosure =
        Box::new(|_| -> CallServiceResult { CallServiceResult::ok(json!([["A"], ["B"]])) });

    let relay_1_peer_id = "relay_1_peer_id";
    let mut relay_1 = create_avm(unit_call_service(), relay_1_peer_id);

    let remote_peer_id = "remove_peer_id";
    let mut remote = create_avm(members_call_service1, remote_peer_id);

    let client_peer_id = "client_peer_id";
    let mut client_1 = create_avm(unit_call_service(), client_peer_id);

    let script = format!(
        r#"
            (seq
                (call "{relay_1_peer_id}" ("identity" "") [] $void1)
                (seq
                    (call "{remote_peer_id}" ("920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9" "get_users") [] members)
                    (fold members m
                        (par
                            (seq
                                (call "{relay_1_peer_id}" ("identity" "") [] $void)
                                (call "{client_peer_id}" ("fgemb3" "add") [m] $void3)
                            )
                            (next m)
                        )
                    )
                )
            )
        "#
    );

    let client_1_result = checked_call_vm!(client_1, <_>::default(), &script, "", "");
    let relay_1_result = checked_call_vm!(relay_1, <_>::default(), &script, client_1_result.data, "");
    let remote_result = checked_call_vm!(remote, <_>::default(), &script, relay_1_result.data, "");
    let relay_1_result = checked_call_vm!(relay_1, <_>::default(), &script, remote_result.data, "");
    let client_1_result = checked_call_vm!(client_1, <_>::default(), script, relay_1_result.data, "");

    let client_1_actual_trace = trace_from_result(&client_1_result);

    let client_1_expected_trace = ExecutionTrace::from(vec![
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        scalar!(
            json!([["A"], ["B"]]),
            peer = remote_peer_id,
            service = "920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9",
            function = "get_users"
        ),
        executed_state::par(2, 3),
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        stream!(
            "result from unit_call_service",
            0,
            peer = client_peer_id,
            service = "fgemb3",
            function = "add",
            args = [json!(["A"])]
        ),
        executed_state::par(2, 0),
        stream!(
            "result from unit_call_service",
            0,
            peer = relay_1_peer_id,
            service = "identity"
        ),
        stream!(
            "result from unit_call_service",
            0,
            peer = client_peer_id,
            service = "fgemb3",
            function = "add",
            args = [json!(["B"])]
        ),
    ]);

    assert_eq!(client_1_actual_trace, client_1_expected_trace);
    assert!(client_1_result.next_peer_pks.is_empty());
}

#[test]
fn init_peer_id() {
    let relay_1_peer_id = "relay_1_peer_id";
    let mut relay_1 = create_avm(unit_call_service(), relay_1_peer_id);

    let client_1_peer_id = "client_1_peer_id";
    let mut client_1 = create_avm(unit_call_service(), client_1_peer_id);

    let initiator_peer_id = "initiator_peer_id";
    let mut initiator = create_avm(unit_call_service(), initiator_peer_id);

    let remote_peer_id = "remote_peer_id";
    let members = json!([[client_1_peer_id], ["B"]]);
    let members_call_service: CallServiceClosure =
        Box::new(move |_| -> CallServiceResult { CallServiceResult::ok(members.clone()) });
    let mut remote = create_avm(members_call_service, remote_peer_id);

    let script = format!(
        r#"(seq
                (seq
                    (call "{relay_1_peer_id}" ("identity" "") [])
                    (seq
                        (call "{remote_peer_id}" ("920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9" "get_users") [] members)
                        (fold members m
                            (par
                                (seq
                                    (call "{relay_1_peer_id}" ("identity" "") [])
                                    (call "{client_1_peer_id}" ("fgemb3" "add") [m])
                                )
                                (next m)
                            )
                        )
                    )
                )
                (call %init_peer_id% ("identity" "") [])
            )
        "#
    );

    let test_params = TestRunParameters::from_init_peer_id(initiator_peer_id);
    let initiator_1_result = checked_call_vm!(initiator, test_params.clone(), &script, "", "");
    let client_1_result = checked_call_vm!(client_1, test_params.clone(), &script, initiator_1_result.data, "");
    let relay_1_result = checked_call_vm!(relay_1, test_params.clone(), &script, client_1_result.data, "");
    let remote_result = checked_call_vm!(remote, test_params.clone(), &script, relay_1_result.data, "");
    let relay_1_result = checked_call_vm!(relay_1, test_params.clone(), &script, remote_result.data, "");
    let client_1_result = checked_call_vm!(client_1, test_params.clone(), &script, relay_1_result.data, "");

    let client_1_actual_trace = trace_from_result(&client_1_result);

    let client_1_expected_trace = ExecutionTrace::from(vec![
        unused!(
            "result from unit_call_service",
            peer = relay_1_peer_id,
            service = "identity"
        ),
        scalar!(
            json!([[client_1_peer_id], ["B"]]),
            peer = remote_peer_id,
            service = "920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9",
            function = "get_users"
        ),
        executed_state::par(2, 3),
        unused!(
            "result from unit_call_service",
            peer = relay_1_peer_id,
            service = "identity"
        ),
        unused!(
            "result from unit_call_service",
            peer = client_1_peer_id,
            service = "fgemb3",
            function = "add",
            args = [json!([client_1_peer_id])]
        ),
        executed_state::par(2, 0),
        unused!(
            "result from unit_call_service",
            peer = relay_1_peer_id,
            service = "identity"
        ),
        unused!(
            "result from unit_call_service",
            peer = client_1_peer_id,
            service = "fgemb3",
            function = "add",
            args = [json!(["B"])]
        ),
        executed_state::request_sent_by(client_1_peer_id),
    ]);

    assert_eq!(client_1_actual_trace, client_1_expected_trace);
    assert_eq!(client_1_result.next_peer_pks, vec![initiator_peer_id.to_string()]);

    let initiator_1_result = checked_call_vm!(initiator, test_params, script, client_1_result.data, "");

    let initiator_1_actual_trace = trace_from_result(&initiator_1_result);

    let initiator_1_expected_trace = ExecutionTrace::from(vec![
        unused!(
            "result from unit_call_service",
            peer = relay_1_peer_id,
            service = "identity"
        ),
        scalar!(
            json!([[client_1_peer_id], ["B"]]),
            peer = remote_peer_id,
            service = "920e3ba3-cbdf-4ae3-8972-0fa2f31fffd9",
            function = "get_users"
        ),
        executed_state::par(2, 3),
        unused!(
            "result from unit_call_service",
            peer = relay_1_peer_id,
            service = "identity"
        ),
        unused!(
            "result from unit_call_service",
            peer = client_1_peer_id,
            service = "fgemb3",
            function = "add",
            args = [json!([client_1_peer_id])]
        ),
        executed_state::par(2, 0),
        unused!(
            "result from unit_call_service",
            peer = relay_1_peer_id,
            service = "identity"
        ),
        unused!(
            "result from unit_call_service",
            peer = client_1_peer_id,
            service = "fgemb3",
            function = "add",
            args = [json!(["B"])]
        ),
        unused!(
            "result from unit_call_service",
            peer = initiator_peer_id,
            service = "identity"
        ),
    ]);

    assert_eq!(initiator_1_actual_trace, initiator_1_expected_trace);
    assert!(initiator_1_result.next_peer_pks.is_empty());
}
