//
// Created by WolverinDEV on 08/03/2020.
//

#include <cassert>
#include <cstring>
#include <netinet/in.h>
#include "ip_router.h"

using namespace ts::network;


constexpr static ip_rounter::route_entry generate_empty_end_node(void*) {
    ip_rounter::route_entry result{};
    for(auto& ptr : result.data)
        ptr = nullptr;
    return result;
}

template <size_t N>
constexpr std::array<ip_rounter::route_entry, N> generate_default_table() noexcept {
    std::array<ip_rounter::route_entry, N> result{};

    for(ip_rounter::route_entry& entry : result)
        entry.use_count = ip_rounter::route_entry::const_flag_mask | 0xFFU;

    for(auto& end_ptr : result[0].data)
        end_ptr = nullptr;

    for(size_t index{1}; index < result.size(); index++)
        for(auto& ptr : result[index].data)
            ptr = &result[index - 1];

    return result;
}

std::array<ip_rounter::route_entry, 16> ip_rounter::recursive_ends = generate_default_table<16>();

struct sockaddr_storage_info {
    size_t address_offset{0};
    size_t address_length{0};
    size_t chunk_offset{0};
};

constexpr std::array<sockaddr_storage_info, 16> generate_storage_info() noexcept {
    std::array<sockaddr_storage_info, 16> result{};
    for(size_t type{0}; type < result.size(); type++) {
        if(type == AF_INET) {
            result[type].address_length = 4;
            result[type].chunk_offset = 12;

            sockaddr_in address{};
            result[type].address_offset = (uintptr_t) &address.sin_addr.s_addr - (uintptr_t) &address;
        } else if(type == AF_INET6) {
            result[type].address_length = 16;
            result[type].chunk_offset = 0;

            sockaddr_in6 address{};
            result[type].address_offset = (uintptr_t) &address.sin6_addr.__in6_u.__u6_addr8 - (uintptr_t) &address;
        }
    }

    return result;
}

std::array<sockaddr_storage_info, 16> storage_infos = generate_storage_info();

inline void address_to_chunks(uint8_t* chunks, const sockaddr_storage &address) {
    if constexpr (AF_INET < 16 && AF_INET6 < 16) {
        /* converter without branches (only one within the memcpy) */
        const auto& info = storage_infos[address.ss_family & 0xFU];

        memset(chunks, 0, 16);
        memcpy(chunks + info.chunk_offset, ((uint8_t*) &address) + info.address_offset, info.address_length); /* we could do this memcpy more performant */
    } else {
        /* converter with branches */
        if(address.ss_family == AF_INET) {
            auto address4 = (sockaddr_in*) &address;

            memset(chunks, 0, 12);
            memcpy(chunks + 12, &address4->sin_addr.s_addr, 4);
        } else if(address.ss_family == AF_INET6) {
            auto address6 = (sockaddr_in6*) &address;
            memcpy(chunks, address6->sin6_addr.__in6_u.__u6_addr8, 16);
        } else {
            memset(chunks, 0, 16);
        }
    }
}

ip_rounter::ip_rounter() {
    for(auto& data : this->root_entry.data)
        data = &ip_rounter::recursive_ends[14];
    this->root_entry.use_count = ip_rounter::route_entry::const_flag_mask | 0xFFU;
}

inline void delete_route_entry(ip_rounter::route_entry* entry, size_t level) {
    if(level != 0) {
        for(auto& data : entry->data) {
            auto e = (ip_rounter::route_entry*) data;
            if(e->is_const_entry()) continue;

            delete_route_entry(e, level - 1);
        }
    }
    delete entry;
}

ip_rounter::~ip_rounter() {
    for(auto& entry : this->unused_nodes)
        delete entry;

    for(auto& data : this->root_entry.data) {
        auto entry = (ip_rounter::route_entry*) data;
        if(entry->is_const_entry()) continue;

        delete_route_entry(entry, 14);
    }
}

/*
 * Because we're only reading memory here, and that even quite fast we do not need to lock the register lock.
 * Even if a block gets changed, it will not be deleted immediately. So we should finish reading first before that memory get freed.
 */
void* ip_rounter::resolve(const sockaddr_storage &address) const {
    uint8_t address_chunks[16];
    address_to_chunks(address_chunks, address);
    const ip_rounter::route_entry* current_chunk = &this->root_entry;

    std::lock_guard lock{this->entry_lock};
    //std::shared_lock lock{this->entry_lock};

#pragma GCC unroll 15
    for(size_t index{0}; index < 15; index++)
        current_chunk = (ip_rounter::route_entry*) current_chunk->data[address_chunks[index]];

    return current_chunk->data[address_chunks[15]];
}

bool ip_rounter::register_route(const sockaddr_storage &address, void *target, void ** old_target) {
    uint8_t address_chunks[16];
    address_to_chunks(address_chunks, address);

    void* _temp_old_target{};
    if(!old_target)
        old_target = &_temp_old_target;

    ip_rounter::route_entry* current_chunk = &this->root_entry;
    std::lock_guard rlock{this->register_lock};

    for(size_t index{0}; index < 15; index++) {
        auto& next_chunk = (ip_rounter::route_entry*&) current_chunk->data[address_chunks[index]];
        if(next_chunk->is_const_entry()) {
            assert(next_chunk == &ip_rounter::recursive_ends[15 - index - 1]);

            auto allocated_entry = this->create_8bit_entry(15 - index - 1);
            if(!allocated_entry) return false;

            next_chunk = allocated_entry;
            current_chunk->use_count++;
        }
        current_chunk = next_chunk;
    }

    *old_target = std::exchange(current_chunk->data[address_chunks[15]], target);
    if(!*old_target) current_chunk->use_count++;
    return true;
}

ip_rounter::route_entry *ip_rounter::create_8bit_entry(size_t level) {
    ip_rounter::route_entry *result;

    if(this->unused_nodes.empty())
        result = new ip_rounter::route_entry{};
    else {
        result = this->unused_nodes.front();
        this->unused_nodes.pop_front();
    }

    result->use_count = 0;
    auto target = level == 0 ? nullptr : &ip_rounter::recursive_ends[level - 1];
    for(auto& data : result->data)
        data = target;

    return result;
}

void *ip_rounter::reset_route(const sockaddr_storage &address) {
    uint8_t address_chunks[16];
    address_to_chunks(address_chunks, address);

    ip_rounter::route_entry* current_chunk = &this->root_entry;
    std::lock_guard rlock{this->register_lock};

    for(size_t index{0}; index < 15; index++) {
        current_chunk = (ip_rounter::route_entry *) current_chunk->data[address_chunks[index]];
        if(current_chunk->is_const_entry()) return nullptr; /* route does not exists */
    }

    auto old = std::exchange(current_chunk->data[address_chunks[15]], nullptr);
    if(!old) return nullptr; /* route does not exists */

    if(--current_chunk->use_count == 0) {
        size_t chunk_index{14};
        do {
            current_chunk = &this->root_entry;
            for(size_t index{0}; index < chunk_index; index++) {
                current_chunk = (ip_rounter::route_entry *) current_chunk->data[address_chunks[index]];
                if(current_chunk->is_const_entry()) return nullptr; /* route does not exists */
            }

            auto& chunk = (ip_rounter::route_entry *&) current_chunk->data[address_chunks[chunk_index]];
            assert(!chunk->is_const_entry());
            assert(chunk->use_count == 0); /* already tested earlier in theory */

            this->unused_nodes.push_back(chunk);
            chunk = &ip_rounter::recursive_ends[15 - chunk_index - 1];

            if(--current_chunk->use_count > 0) break;
        } while(--chunk_index > 0);
    }

    return old;
}

void ip_rounter::cleanup_cache() {
    std::lock_guard rlock{this->register_lock};
    for(auto node : this->unused_nodes)
        delete node;
    this->unused_nodes.clear();
}

bool ip_rounter::validate_chunk_entry(const ip_rounter::route_entry* current_entry, size_t level) const {
    if(current_entry->is_const_entry() && level != 15) /* level 15 is the default root node which is const as well */
        return current_entry == &ip_rounter::recursive_ends[level];

    if(level == 0) return true;

    auto default_pointer = &ip_rounter::recursive_ends[level - 1];
    for(const auto& data_ptr : current_entry->data) {
        if(data_ptr == default_pointer) continue;
        if(!this->validate_chunk_entry((const ip_rounter::route_entry*) data_ptr, level - 1))
            return false;
    }

    return true;
}

bool ip_rounter::validate_tree() const {
    std::lock_guard rlock{this->register_lock};

    /* first lets validate all const chunks */
    for(size_t index{0}; index < 16; index++) {
        auto expected_pointer = index == 0 ? nullptr : &ip_rounter::recursive_ends[index - 1];

        if(!ip_rounter::recursive_ends[index].is_const_entry())
            return false;

        for(const auto& data_ptr : ip_rounter::recursive_ends[index].data)
            if(data_ptr != expected_pointer)
                return false;
    }

    /* not lets check our tree */
    return this->validate_chunk_entry(&this->root_entry, 15);
}

size_t ip_rounter::chunk_memory(const ip_rounter::route_entry *current_entry, size_t level) const {
    size_t result{sizeof(ip_rounter::route_entry)};

    if(level > 0) {
        for(const auto& data_ptr : current_entry->data) {
            auto entry = (const ip_rounter::route_entry*) data_ptr;
            if(entry->is_const_entry()) continue;

            result += chunk_memory(entry, level - 1);
        }
    }

    return result;
}

size_t ip_rounter::used_memory() const {
    return this->chunk_memory(&this->root_entry, 15);
}