/*
    SPDX-FileCopyrightText: 2005 Joris Guisson <joris.guisson@gmail.com>

    SPDX-License-Identifier: GPL-2.0-or-later
*/
#include "udptrackersocket.h"
#include <KLocalizedString>
#include <QHostAddress>
#include <cstdlib>
#include <ctime>
#include <net/portlist.h>
#include <net/serversocket.h>
#include <net/socket.h>
#include <torrent/globals.h>
#include <util/array.h>
#include <util/functions.h>
#include <util/log.h>

#ifdef ERROR
#undef ERROR
#endif

using namespace Qt::Literals::StringLiterals;

namespace bt
{
Uint16 UDPTrackerSocket::port = 4444;

class UDPTrackerSocket::Private : public net::ServerSocket::DataHandler
{
public:
    Private(UDPTrackerSocket *p)
        : p(p)
    {
    }

    ~Private() override
    {
    }

    void listen(const QString &ip, Uint16 port)
    {
        auto sock = std::make_unique<net::ServerSocket>(this);
        if (sock->bind(ip, port)) {
            sockets.emplace_back(std::move(sock));
        }
    }

    bool send(const Uint8 *buf, int size, const net::Address &addr)
    {
        for (const net::ServerSocket::Ptr &sock : std::as_const(sockets)) {
            if (sock->sendTo(buf, size, addr) == size) {
                return true;
            }
        }

        return false;
    }

    void dataReceived(std::unique_ptr<bt::Buffer> buffer, const net::Address &addr) override
    {
        Q_UNUSED(addr);
        if (buffer->size() < 4) {
            return;
        }

        const Uint32 type = ReadUint32(buffer->get(), 0);
        switch (type) {
        case CONNECT:
            p->handleConnect(*buffer);
            break;
        case ANNOUNCE:
            p->handleAnnounce(*buffer);
            break;
        case ERROR:
            p->handleError(*buffer);
            break;
        case SCRAPE:
            p->handleScrape(*buffer);
            break;
        }
    }

    void readyToWrite(net::ServerSocket *sock) override
    {
        Q_UNUSED(sock);
    }

    std::vector<net::ServerSocket::Ptr> sockets;
    QMap<Int32, Action> transactions;
    UDPTrackerSocket *p;
};

UDPTrackerSocket::UDPTrackerSocket()
    : d(std::make_unique<Private>(this))
{
    if (port == 0) {
        port = 4444;
    }

    const QStringList ips = NetworkInterfaceIPAddresses(NetworkInterface());
    for (const QString &ip : ips) {
        d->listen(ip, port);
    }

    if (d->sockets.empty()) {
        // Try all addresses if the previous listen calls all failed
        d->listen(QHostAddress(QHostAddress::AnyIPv6).toString(), port);
        d->listen(QHostAddress(QHostAddress::Any).toString(), port);
    }

    if (d->sockets.empty()) {
        Out(SYS_TRK | LOG_IMPORTANT) << u"Cannot bind to udp port %1"_s.arg(port) << endl;
    } else {
        Globals::instance().getPortList().addNewPort(port, net::UDP, true);
    }
}

UDPTrackerSocket::~UDPTrackerSocket()
{
    Globals::instance().getPortList().removePort(port, net::UDP);
}

void UDPTrackerSocket::sendConnect(Int32 tid, const net::Address &addr)
{
    const Int64 cid = 0x41727101980LL;
    Uint8 buf[16];

    WriteInt64(buf, 0, cid);
    WriteInt32(buf, 8, CONNECT);
    WriteInt32(buf, 12, tid);

    d->send(buf, 16, addr);
    d->transactions.insert(tid, CONNECT);
}

void UDPTrackerSocket::sendAnnounce(Int32 tid, const Uint8 *data, const net::Address &addr)
{
    d->send(data, 98, addr);
    d->transactions.insert(tid, ANNOUNCE);
}

void UDPTrackerSocket::sendScrape(Int32 tid, const bt::Uint8 *data, const net::Address &addr)
{
    d->send(data, 36, addr);
    d->transactions.insert(tid, SCRAPE);
}

void UDPTrackerSocket::cancelTransaction(Int32 tid)
{
    d->transactions.remove(tid);
}

void UDPTrackerSocket::handleConnect(const bt::Buffer &buf)
{
    if (buf.size() < 16) {
        return;
    }

    // Read the transaction_id and check it
    const Int32 tid = ReadInt32(buf.get(), 4);
    const QMap<Int32, Action>::iterator i = d->transactions.find(tid);
    // if we can't find the transaction, just return
    if (i == d->transactions.end()) {
        return;
    }

    // check whether the transaction is a CONNECT
    if (i.value() != CONNECT) {
        d->transactions.erase(i);
        Q_EMIT error(tid, QString());
        return;
    }

    // everything ok, emit signal
    d->transactions.erase(i);
    Q_EMIT connectReceived(tid, ReadInt64(buf.get(), 8));
}

void UDPTrackerSocket::handleAnnounce(const bt::Buffer &buf)
{
    if (buf.size() < 20) {
        return;
    }

    // Read the transaction_id and check it
    const Int32 tid = ReadInt32(buf.get(), 4);
    const QMap<Int32, Action>::iterator i = d->transactions.find(tid);
    // if we can't find the transaction, just return
    if (i == d->transactions.end()) {
        return;
    }

    // check whether the transaction is a ANNOUNCE
    if (i.value() != ANNOUNCE) {
        d->transactions.erase(i);
        Q_EMIT error(tid, QString());
        return;
    }

    // everything ok, emit signal
    d->transactions.erase(i);
    Q_EMIT announceReceived(tid, buf.get(), buf.size());
}

void UDPTrackerSocket::handleError(const bt::Buffer &buf)
{
    if (buf.size() < 8) {
        return;
    }

    // Read the transaction_id and check it
    const Int32 tid = ReadInt32(buf.get(), 4);
    const QMap<Int32, Action>::iterator it = d->transactions.find(tid);
    // if we can't find the transaction, just return
    if (it == d->transactions.end()) {
        return;
    }

    // extract error message
    d->transactions.erase(it);
    QString msg;
    for (Uint32 i = 8; i < buf.size(); i++) {
        msg += QLatin1Char(buf.get()[i]);
    }

    // emit signal
    Q_EMIT error(tid, msg);
}

void UDPTrackerSocket::handleScrape(const bt::Buffer &buf)
{
    if (buf.size() < 20) {
        return;
    }

    // Read the transaction_id and check it
    const Int32 tid = ReadInt32(buf.get(), 4);
    const QMap<Int32, Action>::iterator i = d->transactions.find(tid);
    // if we can't find the transaction, just return
    if (i == d->transactions.end()) {
        return;
    }

    // check whether the transaction is a SCRAPE
    if (i.value() != SCRAPE) {
        d->transactions.erase(i);
        Q_EMIT error(tid, QString());
        return;
    }

    // everything ok, emit signal
    d->transactions.erase(i);
    Q_EMIT scrapeReceived(tid, buf.get(), buf.size());
}

Int32 UDPTrackerSocket::newTransactionID()
{
    Int32 transaction_id = rand() * time(nullptr);
    while (d->transactions.contains(transaction_id)) {
        transaction_id++;
    }
    return transaction_id;
}

void UDPTrackerSocket::setPort(Uint16 p)
{
    port = p;
}

Uint16 UDPTrackerSocket::getPort()
{
    return port;
}
}

#include "moc_udptrackersocket.cpp"
