#include <test/csf/BasicNetwork.h>
#include <test/csf/Scheduler.h>

#include <xrpl/beast/unit_test.h>

#include <set>
#include <vector>

namespace xrpl {
namespace test {

class BasicNetwork_test : public beast::unit_test::suite
{
public:
    struct Peer
    {
        int id;
        std::set<int> set;

        Peer(Peer const&) = default;
        Peer(Peer&&) = default;

        explicit Peer(int id_) : id(id_)
        {
        }

        template <class Net>
        void
        start(csf::Scheduler& scheduler, Net& net)
        {
            using namespace std::chrono_literals;
            auto t = scheduler.in(1s, [&] { set.insert(0); });
            if (id == 0)
            {
                for (auto const link : net.links(this))
                    net.send(this, link.target, [&, to = link.target] {
                        to->receive(net, this, 1);
                    });
            }
            else
            {
                scheduler.cancel(t);
            }
        }

        template <class Net>
        void
        receive(Net& net, Peer* from, int m)
        {
            set.insert(m);
            ++m;
            if (m < 5)
            {
                for (auto const link : net.links(this))
                    net.send(this, link.target, [&, mm = m, to = link.target] {
                        to->receive(net, this, mm);
                    });
            }
        }
    };

    void
    testNetwork()
    {
        using namespace std::chrono_literals;
        std::vector<Peer> pv;
        pv.emplace_back(0);
        pv.emplace_back(1);
        pv.emplace_back(2);
        csf::Scheduler scheduler;
        csf::BasicNetwork<Peer*> net(scheduler);
        BEAST_EXPECT(!net.connect(&pv[0], &pv[0]));
        BEAST_EXPECT(net.connect(&pv[0], &pv[1], 1s));
        BEAST_EXPECT(net.connect(&pv[1], &pv[2], 1s));
        BEAST_EXPECT(!net.connect(&pv[0], &pv[1]));
        for (auto& peer : pv)
            peer.start(scheduler, net);
        BEAST_EXPECT(scheduler.step_for(0s));
        BEAST_EXPECT(scheduler.step_for(1s));
        BEAST_EXPECT(scheduler.step());
        BEAST_EXPECT(!scheduler.step());
        BEAST_EXPECT(!scheduler.step_for(1s));
        net.send(&pv[0], &pv[1], [] {});
        net.send(&pv[1], &pv[0], [] {});
        BEAST_EXPECT(net.disconnect(&pv[0], &pv[1]));
        BEAST_EXPECT(!net.disconnect(&pv[0], &pv[1]));
        for (;;)
        {
            auto const links = net.links(&pv[1]);
            if (links.empty())
                break;
            BEAST_EXPECT(net.disconnect(&pv[1], links[0].target));
        }
        BEAST_EXPECT(pv[0].set == std::set<int>({0, 2, 4}));
        BEAST_EXPECT(pv[1].set == std::set<int>({1, 3}));
        BEAST_EXPECT(pv[2].set == std::set<int>({2, 4}));
    }

    void
    testDisconnect()
    {
        using namespace std::chrono_literals;
        csf::Scheduler scheduler;
        csf::BasicNetwork<int> net(scheduler);
        BEAST_EXPECT(net.connect(0, 1, 1s));
        BEAST_EXPECT(net.connect(0, 2, 2s));

        std::set<int> delivered;
        net.send(0, 1, [&]() { delivered.insert(1); });
        net.send(0, 2, [&]() { delivered.insert(2); });

        scheduler.in(1000ms, [&]() { BEAST_EXPECT(net.disconnect(0, 2)); });
        scheduler.in(1100ms, [&]() { BEAST_EXPECT(net.connect(0, 2)); });

        scheduler.step();

        // only the first message is delivered because the disconnect at 1 s
        // purges all pending messages from 0 to 2
        BEAST_EXPECT(delivered == std::set<int>({1}));
    }

    void
    run() override
    {
        testNetwork();
        testDisconnect();
    }
};

BEAST_DEFINE_TESTSUITE(BasicNetwork, csf, xrpl);

}  // namespace test
}  // namespace xrpl