#include "../include/sibs/BootstrapConnection.hpp"
#include "../include/Log.hpp"
#include <sibs/SafeDeserializer.hpp>

namespace sibs
{
    BootstrapConnection::BootstrapConnection(const Ipv4 &bootstrapAddress)
    {
        serverPeer = connections.connectServer(bootstrapAddress, std::bind(&BootstrapConnection::receiveDataFromServer, this, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3));
    }
    
    // TODO: This is vulnerable against MitM attack, replace with asymmetric cryptography, get data signed with server private key and verify against known server public key
    void BootstrapConnection::receiveDataFromServer(std::shared_ptr<DirectConnectionPeer> peer, const void *data, const usize size)
    {
        sibs::SafeDeserializer deserializer((const u8*)data, size);
        PubsubKey pubsubKey;
        deserializer.extract(pubsubKey.data.data(), PUBSUB_KEY_LENGTH);
        
        listenerCallbackFuncMutex.lock();
        auto listenerFuncIt = listenCallbackFuncs.find(pubsubKey);
        if(listenerFuncIt == listenCallbackFuncs.end())
        {
            Log::debug("BoostrapConnection: No listener found for key XXX, ignoring...");
            listenerCallbackFuncMutex.unlock();
            return;
        }
        auto listenerCallbackFunc = listenerFuncIt->second;
        listenerCallbackFuncMutex.unlock();
        
        auto &peers = subscribedPeers[pubsubKey];
        while(!deserializer.empty())
        {
            sa_family_t addressFamily = deserializer.extract<u32>();
            if(addressFamily == AF_INET)
            {
                in_addr_t ipv4Address = deserializer.extract<u32>();
                u16 port = deserializer.extract<u32>();
                Ipv4 newPeerAddress;
                newPeerAddress.address.sin_family = addressFamily;
                newPeerAddress.address.sin_addr.s_addr = ipv4Address;
                newPeerAddress.address.sin_port = port;
                memset(newPeerAddress.address.sin_zero, 0, sizeof(newPeerAddress.address.sin_zero));
                // TODO: Move connection to thread and add callback function, just like @receiveData and @send
                std::shared_ptr<DirectConnectionPeer> newPeer = connections.connect(newPeerAddress, std::bind(&BootstrapConnection::receiveDataFromPeer, this, listenerCallbackFunc, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3));
                peers.push_back(newPeer);
            }
            else
                Log::error("BootstrapConnection: Unknown address family: %d", addressFamily);
        }
    }
    
    void BootstrapConnection::receiveDataFromPeer(BoostrapConnectionListenCallbackFunc listenCallbackFunc, std::shared_ptr<DirectConnectionPeer> peer, const void *data, const usize size)
    {
        if(listenCallbackFunc)
            listenCallbackFunc(data, size);
    }
    
    void BootstrapConnection::listen(const PubsubKey &pubsubKey, BoostrapConnectionListenCallbackFunc callbackFunc)
    {
        {
            std::lock_guard<std::mutex> lock(listenerCallbackFuncMutex);
            if(listenCallbackFuncs.find(pubsubKey) != listenCallbackFuncs.end())
                throw PubsubKeyAlreadyListeningException("");
            listenCallbackFuncs[pubsubKey] = callbackFunc;
        }
        connections.send(serverPeer, std::make_shared<std::vector<u8>>(pubsubKey.data.data(), pubsubKey.data.data()));
    }
    
    void BootstrapConnection::put(const PubsubKey &pubsubKey, std::shared_ptr<std::vector<u8>> data)
    {
        {
            std::lock_guard<std::mutex> lock(listenerCallbackFuncMutex);
            auto listenCallbackFuncIt = listenCallbackFuncs.find(pubsubKey);
            if(listenCallbackFuncIt != listenCallbackFuncs.end() && listenCallbackFuncIt->second)
                listenCallbackFuncIt->second(data->data(), data->size());
        }
        
        auto peersIt = subscribedPeers.find(pubsubKey);
        if(peersIt == subscribedPeers.end())
            return;
        
        for(auto &peer : peersIt->second)
        {
            connections.send(peer, data);
        }
    }
}