/// @file dpf/net/memory_sink.hpp
/// @brief In-process paired RoundSink for correctness tests.
#ifndef LIBDPF_INCLUDE_DPF_NET_MEMORY_SINK_HPP__
#define LIBDPF_INCLUDE_DPF_NET_MEMORY_SINK_HPP__

#include <cstddef>
#include <cstdint>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <utility>
#include <vector>

#include "dpf/net/round_sink.hpp"

namespace dpf
{
namespace net
{

/// @brief Shared state between the two ends of a memory sink pair.
struct memory_sink_hub
{
    std::size_t count = 0;
    std::vector<std::size_t> slot_bytes;
    std::vector<round_window> a;
    std::vector<round_window> b;
    mutable std::mutex mu;

    explicit memory_sink_hub(std::size_t n, std::vector<std::size_t> slots)
      : count(n), slot_bytes(std::move(slots))
    {
        a.reserve(slot_bytes.size());
        b.reserve(slot_bytes.size());
        for (std::size_t sb : slot_bytes)
        {
            a.emplace_back(count, sb);
            b.emplace_back(count, sb);
        }
    }
};

/// @brief One end of a memory-paired RoundSink.
class memory_sink : public RoundSink
{
public:
    memory_sink(std::shared_ptr<memory_sink_hub> hub, bool side_a)
      : hub_(std::move(hub)), side_a_(side_a)
    {
        if (!hub_)
            throw std::invalid_argument("memory_sink needs a hub");
    }

    std::size_t count() const noexcept override { return hub_->count; }
    std::size_t rounds() const noexcept override
    {
        return hub_->slot_bytes.size();
    }

    std::size_t slot_bytes(std::uint16_t round) const override
    {
        if (round >= hub_->slot_bytes.size())
            throw std::out_of_range("memory_sink round");
        return hub_->slot_bytes[round];
    }

    void submit(std::uint16_t round, std::size_t index,
        const std::uint8_t * bytes, std::size_t n) override
    {
        std::lock_guard<std::mutex> lock(hub_->mu);
        mine(round).submit(index, bytes, n);
    }

    bool peer_ready(std::uint16_t round, std::size_t index) const override
    {
        std::lock_guard<std::mutex> lock(hub_->mu);
        return mine(round).peer_ready(index);
    }

    void read_peer(std::uint16_t round, std::size_t index, std::uint8_t * out,
        std::size_t n) const override
    {
        std::lock_guard<std::mutex> lock(hub_->mu);
        mine(round).read_peer(index, out, n);
    }

    void flush() override
    {
        std::lock_guard<std::mutex> lock(hub_->mu);
        flush_unlocked();
    }

    void flush_round(std::uint16_t round) override
    {
        std::lock_guard<std::mutex> lock(hub_->mu);
        flush_round_unlocked(round);
    }

    void poll() override {}

private:
    void flush_unlocked()
    {
        for (std::uint16_t r = 0; r < hub_->slot_bytes.size(); ++r)
        {
            std::size_t begin = 0;
            std::size_t n = 0;
            mine(r).pending_out(begin, n);
            if (n != 0)
            {
                flush_round_unlocked(r);
                continue;
            }
            peer(r).pending_out(begin, n);
            if (n != 0)
                flush_round_unlocked(r);
        }
    }

    void flush_round_unlocked(std::uint16_t round)
    {
        if (round >= hub_->slot_bytes.size())
            throw std::out_of_range("memory_sink flush_round");
        auto & local = mine(round);
        auto & remote = peer(round);
        std::size_t begin_l = 0;
        std::size_t n_l = 0;
        const std::uint8_t * pend_l = local.pending_out(begin_l, n_l);
        std::size_t begin_r = 0;
        std::size_t n_r = 0;
        const std::uint8_t * pend_r = remote.pending_out(begin_r, n_r);
        if (n_l != 0)
        {
            remote.accept_peer_at(begin_l, pend_l, n_l);
            local.mark_flushed(n_l);
        }
        if (n_r != 0)
        {
            local.accept_peer_at(begin_r, pend_r, n_r);
            remote.mark_flushed(n_r);
        }
    }
    round_window & mine(std::uint16_t round)
    {
        if (round >= hub_->slot_bytes.size())
            throw std::out_of_range("memory_sink round");
        return side_a_ ? hub_->a[round] : hub_->b[round];
    }

    const round_window & mine(std::uint16_t round) const
    {
        if (round >= hub_->slot_bytes.size())
            throw std::out_of_range("memory_sink round");
        return side_a_ ? hub_->a[round] : hub_->b[round];
    }

    round_window & peer(std::uint16_t round)
    {
        if (round >= hub_->slot_bytes.size())
            throw std::out_of_range("memory_sink round");
        return side_a_ ? hub_->b[round] : hub_->a[round];
    }

    std::shared_ptr<memory_sink_hub> hub_;
    bool side_a_ = true;
};

/// @brief Build a connected pair of memory sinks that share one hub.
inline std::pair<memory_sink, memory_sink> make_memory_sink_pair(
    std::size_t count, std::vector<std::size_t> slot_bytes)
{
    auto hub = std::make_shared<memory_sink_hub>(count, std::move(slot_bytes));
    return {memory_sink(hub, true), memory_sink(hub, false)};
}

}  // namespace net
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_NET_MEMORY_SINK_HPP__