#include "network/switch/multicast_table.h"

#include <set>
#include <iterator>

namespace dofs {

MulticastTable::MulticastTable()
        : _groups(MAX_GROUPS) {}

std::size_t MulticastTable::group_count() const noexcept {
        return _groups.size();
}

bool MulticastTable::add_tree(std::size_t group_id, uint16_t tree_id) {
        if (group_id >= _groups.size())
                return false;

        auto& trees = _groups[group_id];
        auto it = std::find_if(trees.begin(), trees.end(),
        [&](const McTree & t) {
                return t.tree_id == tree_id;
        });

        if (it != trees.end())
                return false;

        McTree t;
        t.tree_id = tree_id;
        trees.insert(std::lower_bound(trees.begin(), trees.end(), tree_id,
        [](const McTree & a, uint16_t id) {
                return a.tree_id < id;
        }),
        t);
        return true;
}

bool MulticastTable::delete_tree(std::size_t group_id, uint16_t tree_id) {
        if (group_id >= _groups.size())
                return false;

        auto& trees = _groups[group_id];
        auto it = std::find_if(trees.begin(), trees.end(),
        [&](const McTree & t) {
                return t.tree_id == tree_id;
        });

        if (it == trees.end())
                return false;

        trees.erase(it);
        return true;
}

bool MulticastTable::add_child_port(std::size_t group_id, uint16_t tree_id,
                                    PortId out_port) {
        if (group_id >= _groups.size())
                return false;

        auto& trees = _groups[group_id];
        auto it = std::find_if(trees.begin(), trees.end(),
        [&](const McTree & t) {
                return t.tree_id == tree_id;
        });

        if (it == trees.end())
                return false;

        return insert_sorted_unique(it->child_ports, out_port);
}

bool MulticastTable::delete_child_port(std::size_t group_id, uint16_t tree_id,
                                       PortId out_port) {
        if (group_id >= _groups.size())
                return false;

        auto& trees = _groups[group_id];
        auto it = std::find_if(trees.begin(), trees.end(),
        [&](const McTree & t) {
                return t.tree_id == tree_id;
        });

        if (it == trees.end())
                return false;

        return erase_sorted(it->child_ports, out_port);
}

bool MulticastTable::set_parent(std::size_t group_id, uint16_t tree_id,
                                PortId parent) {
        if (group_id >= _groups.size())
                return false;

        auto& trees = _groups[group_id];
        auto it = std::find_if(trees.begin(), trees.end(),
        [&](const McTree & t) {
                return t.tree_id == tree_id;
        });

        if (it == trees.end())
                return false;

        it->parent_port = parent;
        return true;
}

bool MulticastTable::set_weight(std::size_t group_id, uint16_t tree_id,
                                uint8_t w) {
        if (group_id >= _groups.size())
                return false;

        auto& trees = _groups[group_id];
        auto it = std::find_if(trees.begin(), trees.end(),
        [&](const McTree & t) {
                return t.tree_id == tree_id;
        });

        if (it == trees.end())
                return false;

        it->weight = w;
        return true;
}

bool MulticastTable::set_epoch(std::size_t group_id, uint16_t tree_id,
                               uint8_t epoch) {
        if (group_id >= _groups.size())
                return false;

        auto& trees = _groups[group_id];
        auto it = std::find_if(trees.begin(), trees.end(),
        [&](const McTree & t) {
                return t.tree_id == tree_id;
        });

        if (it == trees.end())
                return false;

        it->epoch = epoch;
        return true;
}

const std::vector<McTree> *MulticastTable::trees_of(std::size_t group_id)
const {
        if (group_id >= _groups.size())
                return nullptr;

        return &_groups[group_id];
}

// Legacy: union of tree_id=0 ports for all set groups
std::vector<PortId> MulticastTable::get_port_list(PacketGroups groups) const {
        std::vector<PortId> result;

        for (std::size_t gid = 0; gid < _groups.size(); ++gid) {
                const bool bit_set = (groups >> gid) & static_cast<PacketGroups>(1);

                if (!bit_set)
                        continue;

                const auto& trees = _groups[gid];
                auto it = std::find_if(trees.begin(), trees.end(),
                [](const McTree & t) {
                        return t.tree_id == 0;
                });

                if (it == trees.end())
                        continue;

                if (it->child_ports.empty())
                        continue;

                if (result.empty()) {
                        result = it->child_ports;
                } else {
                        std::vector<PortId> merged;
                        merged.reserve(result.size() + it->child_ports.size());
                        std::set_union(result.begin(), result.end(),
                                       it->child_ports.begin(), it->child_ports.end(),
                                       std::back_inserter(merged));
                        result.swap(merged);
                }
        }

        return result;
}

bool MulticastTable::insert_sorted_unique(std::vector<PortId> &vec, PortId v) {
        auto it = std::lower_bound(vec.begin(), vec.end(), v);

        if (it != vec.end() && *it == v)
                return false;

        vec.insert(it, v);
        return true;
}

bool MulticastTable::erase_sorted(std::vector<PortId> &vec, PortId v) {
        auto it = std::lower_bound(vec.begin(), vec.end(), v);

        if (it == vec.end() || *it != v)
                return false;

        vec.erase(it);
        return true;
}

} // namespace dofs
