/*
 * This file is part of the Flowee project
 * Copyright (C) 2020-2023 Tom Zander <tom@flowee.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "PeerAddressDB.h"
#include "ConnectionManager.h"
#include <Message.h>
#include <streaming/BufferPool.h>
#include <streaming/MessageBuilder.h>
#include <streaming/MessageParser.h>
#include <streaming/P2PParser.h>
#include <random.h>

#include <fstream>

enum SavingTags {
    Separator = 0,
    Hostname,
    IPAddress,
    Port,
    Services,
    LastConnected,
    Punishment,
    Segment,
    EverConnected,
    LastReceivedGoodHeaders
    // AskedAddr?
};

const EndPoint &PeerAddress::peerAddress() const
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    return i->second.address;
}

void PeerAddress::successfullyConnected()
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    time_t now = time(NULL);
    assert(now > 0);
    i->second.lastConnected = static_cast<uint32_t>(now);
    if (i->second.punishment > 500)
        i->second.punishment -= 125;
    else if (i->second.punishment < 20) // the 'never connected' punishment.
        i->second.punishment = 0;
    i->second.inUse = true;
    i->second.everConnected = true;
}

void PeerAddress::gotGoodHeaders()
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    time_t now = time(nullptr);
    assert(now > 0);
    i->second.lastConnected = static_cast<uint32_t>(now);
    if (i->second.punishment > 500)
        i->second.punishment -= 270;
    i->second.lastReceivedGoodHeaders = now;
}

short PeerAddress::punishPeer(short amount)
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    int newPunishment = int(i->second.punishment) + amount;
    if (i->second.punishment < PUNISHMENT_MAX && newPunishment >= PUNISHMENT_MAX) {
        d->m_disabledPeerCount++;
    } else if (i->second.punishment >= PUNISHMENT_MAX && newPunishment < PUNISHMENT_MAX) {
        d->m_disabledPeerCount--;
    }

    i->second.punishment = std::min(newPunishment, 0xeFFF); // avoid overflow
    return i->second.punishment;
}

short PeerAddress::punishment() const
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    return i->second.punishment;
}

void PeerAddress::resetPunishment()
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    i->second.punishment = 0;
    // update lastConnected to now without setting 'everConnected' to
    // at least remember when we reset the punishment.
    i->second.lastConnected = static_cast<uint32_t>(time(nullptr));
}

bool PeerAddress::isValid() const
{
    return d && m_id >= 0 && d->m_nextPeerId > m_id;
}

bool PeerAddress::askedAddresses() const
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    return i->second.askedAddr;
}

void PeerAddress::setAskedAddresses(bool on)
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    i->second.askedAddr = on;
}

bool PeerAddress::hasEverConnected() const
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    return i->second.everConnected;
}

uint32_t PeerAddress::lastReceivedGoodHeaders() const
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    return i->second.lastReceivedGoodHeaders;
}

uint16_t PeerAddress::segment() const
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    return i->second.segment;
}

void PeerAddress::setSegment(uint16_t segment)
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    i->second.segment = segment;
}

void PeerAddress::setInUse(bool on)
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    i->second.inUse = on;
}

void PeerAddress::setServices(uint64_t services)
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    i->second.services = services;
}

uint32_t PeerAddress::lastConnected() const
{
    auto i = d->m_peers.find(m_id);
    assert(d->m_peers.end() != i);
    return i->second.lastConnected;
}

PeerAddress::PeerAddress(PeerAddressDB *parent, int peerId)
    : d(parent), m_id(peerId)
{
    assert(parent);
}

int PeerAddress::id() const
{
    return m_id;
}


// ////////////////////////////////////////

PeerAddressDB::PeerAddressDB(ConnectionManager *parent)
    : m_parent(parent)
{
}

PeerAddress PeerAddressDB::findBest(uint64_t requiredServices, uint16_t segment)
{
    std::unique_lock<std::mutex> lock(m_lock);
    if (m_nextPeerId == 0)
        return PeerAddress(this, -1);

    int usefulCount = 0;
    int best = -1;
    int bestScore = 0;
    for (int attempts = 0; attempts < 1000 && usefulCount < 100; ++attempts) {
        const int i = GetRandInt(m_nextPeerId - 1);
        assert(i >= 0);
        assert(m_peers.size() > static_cast<size_t>(i));
        const PeerInfo &info = m_peers[i];
        if (info.inUse)
            continue;
        if ((info.services & requiredServices) != requiredServices)
            continue;
        if (segment != 0 && info.segment != 0 && segment != info.segment)
            continue;
        if (!m_supportIPv4Net && info.address.ipAddress.is_v4())
            continue;
        if (!m_supportIPv6Net && info.address.ipAddress.is_v6())
            continue;

        ++usefulCount;
        const int hoursAgoConnected = (time(nullptr) - info.lastConnected) / 3600;
        /*
         * A punishment-score (aka ban-score) isn't everything.
         * If its been a long time since we connected, maybe the peer
         * has changed to be better.
         */
        int punishment = info.punishment;
        if (info.everConnected && hoursAgoConnected > 7 * 24 * 6)
            punishment /= 2;
        int score = std::max(0, 1000 - punishment);

        // If its all the same, I prefer nodes that we recently connected to (3 weeks)
        if (info.everConnected)
            score += 250 - std::min(250, hoursAgoConnected / 2);

        if (info.address.announcePort != m_defaultPortNr) // SLIGHTLY prefer non-default port
            score += 50;

        if (score > bestScore) {
            bestScore = score;
            best = i;
        }
    }
    return PeerAddress(this, best);
}

int PeerAddressDB::peerCount() const
{
    std::unique_lock<std::mutex> lock(m_lock);
    return static_cast<int>(m_peers.size()) - m_disabledPeerCount;
}

void PeerAddressDB::processAddressMessage(const Message &message, int sourcePeerId)
{
    std::unique_lock<std::mutex> lock(m_lock);
    // TODO Should we ignore messages here coming from peers that are not on the same chain?
    size_t oldCount = m_peers.size();
    try {
        Streaming::P2PParser parser(message);
        const size_t count = parser.readCompactInt();
        logInfo() << "Received" << count << "addresses" << "from peer:" << sourcePeerId;
        for (size_t i = 0; i < count; ++i) {
            PeerInfo info;
            info.lastConnected = parser.readInt();
            info.services = parser.readLong();
            auto ip = parser.readBytes(16);
            auto port = parser.readWordBE();
            info.address = EndPoint::fromAddr(ip, port);
            insert(info);
        }
    } catch (const std::runtime_error &err) {
        logInfo() << "Failed to read address message from peer:" << sourcePeerId;
        m_parent->punish(sourcePeerId);
        return;
    }
    if (oldCount != m_peers.size()) {
        logInfo().nospace() << "We now have " << m_peers.size() << " addresses (thanks! peer: " << sourcePeerId << ")";
        m_needsSave = true;
    }
}

void PeerAddressDB::addOne(const EndPoint &endPoint)
{
    std::unique_lock<std::mutex> lock(m_lock);
    PeerInfo info;
    info.address = endPoint;
    info.services = 5;
    insert(info);
    m_needsSave = true;
}

void PeerAddressDB::saveDatabase(const boost::filesystem::path &basedir)
{
    std::unique_lock<std::mutex> lock(m_lock);
    if (!m_needsSave)
        return;
    Streaming::BufferPool pool(m_peers.size() * 40);
    Streaming::MessageBuilder builder(pool);
    char ip[16];
    for (const auto &item : m_peers) {
        if (item.second.address.ipAddress.is_unspecified()) {
            builder.add(Hostname, item.second.address.hostname);
        } else {
            item.second.address.toAddr(ip);
            builder.addByteArray(IPAddress, ip, 16);
        }
        if (item.second.address.announcePort != m_defaultPortNr)
            builder.add(Port, item.second.address.announcePort);
        builder.add(Services, item.second.services);
        builder.add(LastConnected, uint64_t(item.second.lastConnected));
        if (item.second.punishment > 0)
            builder.add(Punishment, item.second.punishment);
        if (item.second.segment != 0)
            builder.add(Segment, item.second.segment);
        if (item.second.everConnected)
            builder.add(EverConnected, true);
        if (item.second.lastReceivedGoodHeaders)
            builder.add(LastReceivedGoodHeaders, uint64_t(item.second.lastReceivedGoodHeaders));

        builder.add(Separator, true); // separator
    }
    auto data = builder.buffer();

    try {
        boost::system::error_code error;
        boost::filesystem::create_directories(basedir, error);
        if (error && !boost::filesystem::exists(basedir) && !boost::filesystem::is_directory(basedir)) {
            logFatal() << "P2P.PeerAddressDB can't save. Failed creating the dir:" << basedir.string();
            return;
        }
        boost::filesystem::remove(basedir / "peers.dat~");

        std::ofstream outFile((basedir / "peers.dat~").string());
        outFile.write(data.begin(), data.size());

        boost::filesystem::rename(basedir / "peers.dat~", basedir / "peers.dat");
        m_needsSave = false;
    } catch (const std::exception &e) {
        logFatal() << "Failed to save the database. Reason:" << e.what();
    }
}

void PeerAddressDB::loadDatabase(const boost::filesystem::path &basedir)
{
    std::unique_lock<std::mutex> lock(m_lock);
    std::ifstream in((basedir / "peers.dat").string());
    if (!in.is_open())
        return;
    m_nextPeerId = 0;
    const auto dataSize = boost::filesystem::file_size(basedir / "peers.dat");
    Streaming::BufferPool pool(dataSize);
    in.read(pool.begin(), dataSize);
    Streaming::MessageParser parser(pool.commit(dataSize));
    PeerInfo info;
    while (parser.next() == Streaming::FoundTag) {
        if (parser.tag() == Separator) {
            if (info.address.isValid())
                m_peers.insert(std::make_pair(m_nextPeerId++, info));
            info = PeerInfo();
            info.everConnected = true; // defaults in saving that differ from struct defaults
        }
        else if (parser.tag() == IPAddress) {
            info.address = EndPoint::fromAddr(parser.bytesData(), m_defaultPortNr);
        }
        else if (parser.tag() == Hostname) {
            info.address = EndPoint(parser.stringData(), m_defaultPortNr);
        }
        else if (parser.tag() == Port) {
            info.address.announcePort = info.address.peerPort = parser.intData();
        }
        else if (parser.tag() == Services) {
            info.services = parser.longData();
        }
        else if (parser.tag() == LastConnected) {
            info.lastConnected = parser.longData();
        }
        else if (parser.tag() == Punishment) {
            info.punishment = parser.intData();
            if (info.punishment >= PUNISHMENT_MAX)
                m_disabledPeerCount++;
        }
        else if (parser.tag() == Segment) {
            info.segment = parser.intData();
        }
        else if (parser.tag() == EverConnected) {
            info.everConnected = parser.boolData();
        }
        else if (parser.tag() == LastReceivedGoodHeaders) {
            info.lastReceivedGoodHeaders = parser.longData();
        }
    }
    m_needsSave = false;

    logInfo() << "Peer database loaded:" << m_peers.size() << "disabled peer count:" << m_disabledPeerCount;
}

void PeerAddressDB::insert(PeerInfo &pi)
{
    if (pi.address.ipAddress.is_unspecified()) // try to see if hostname is an IP. If so, bypass DNS lookup
        try { pi.address.ipAddress = boost::asio::ip::address::from_string(pi.address.hostname); } catch (...) {}
    const bool hasIp = !pi.address.ipAddress.is_unspecified();
    if (!hasIp && pi.address.hostname.empty())
        return;
    for (auto i = m_peers.begin(); i != m_peers.end(); ++i) {
        if (hasIp && i->second.address.ipAddress == pi.address.ipAddress) {
            if (i->second.punishment > 400)
                i->second.punishment -= 100;
            return;
        }
        if (!hasIp && i->second.address.hostname == pi.address.hostname)
            return;
    }

    m_peers.insert(std::make_pair(m_nextPeerId++, pi));
    m_needsSave = true;
}

int PeerAddressDB::defaultPortNr() const
{
    return m_defaultPortNr;
}

void PeerAddressDB::setDefaultPortNr(int defaultPortNr)
{
    m_defaultPortNr = defaultPortNr;
}