/*
 * Copyright (C) 2010  Barracuda Networks, Inc.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this library.  If not, see <https://www.gnu.org/licenses/>.
 *
 */

#include "udpportreserver.h"

#include <QUdpSocket>
#include <stdlib.h>

namespace XMPP {
class UdpPortReserver::Private : public QObject {
    Q_OBJECT

public:
    class Item {
    public:
        int  port; // port to reserve
        bool lent;

        // list of sockets for this port, one socket per address.
        //   note that we may have sockets bound for addresses
        //   we no longer care about, if we are currently lending
        //   them out
        QList<QUdpSocket *> sockList;

        // keep track of which addresses we lent out
        QList<QHostAddress> lentAddrs;

        Item() : port(-1), lent(false) { }

        bool haveAddress(const QHostAddress &addr) const
        {
            for (const QUdpSocket *sock : sockList) {
                if (sock->localAddress() == addr)
                    return true;
            }

            return false;
        }
    };

    UdpPortReserver *   q;
    QList<QHostAddress> addrs;
    QList<int>          ports; // sorted.

    // addrs * ports = all available sockets

    /**
     * @brief items kind of ports slice over all provided addresses.
     * When we request binding on one port it means it will be bound on one item (all its addresses).
     * The items are sorted by port.
     */
    QList<Item> items;

    Private(UdpPortReserver *_q) : QObject(_q), q(_q) { }

    ~Private()
    {

        bool lendingAny = std::any_of(items.begin(), items.end(), [](auto const &i) { return i.lent; });

        Q_ASSERT(!lendingAny);
        if (lendingAny)
            abort();

        for (const Item &i : qAsConst(items)) {
            for (QUdpSocket *sock : i.sockList)
                sock->deleteLater();
        }
    }

    void updateAddresses(const QList<QHostAddress> &newAddrs)
    {
        addrs = newAddrs;

        tryBind();
        tryCleanup();
    }

    void updatePorts(const QList<int> &newPorts)
    {
        QList<int> added;
        for (int x : newPorts) {
            bool found = false;
            for (const Item &i : qAsConst(items)) {
                if (i.port == x) {
                    found = true;
                    break;
                }
            }

            if (!found)
                added += x;
        }

        ports = newPorts;

        // keep ports in sorted order
        std::sort(ports.begin(), ports.end());

        for (int x : qAsConst(added)) {
            int insert_before = items.count();
            for (int n = 0; n < items.count(); ++n) {
                if (x < items[n].port) {
                    insert_before = n;
                    break;
                }
            }

            Item i;
            i.port = x;
            items.insert(insert_before, i);
        }

        tryBind();
        tryCleanup();
    }

    bool reservedAll() const
    {
        bool ok = true;
        for (const Item &i : items) {
            // skip ports we don't care about
            if (!ports.contains(i.port))
                continue;

            if (!isReserved(i)) {
                ok = false;
                break;
            }
        }

        return ok;
    }

    QList<QUdpSocket *> borrowSockets(int portCount, QObject *parent)
    {
        Q_ASSERT(portCount > 0);

        QList<QUdpSocket *> out;

        if (portCount > 1) {
            // first try to see if we can find something all in a
            //   row, starting with best alignment to worst
            for (int align = portCount; align >= 2; align /= 2) {
                int at = findConsecutive(portCount, align);
                if (at != -1) {
                    for (int n = 0; n < portCount; ++n)
                        out += lendItem(&items[at + n], parent);

                    break;
                }
            }

            if (out.isEmpty()) {
                // otherwise, try splitting them up into
                //   smaller consecutive chunks
                int chunks[2];
                chunks[0] = portCount / 2 + (portCount % 2);
                chunks[1] = portCount / 2;
                for (int n = 0; n < 2; ++n)
                    out += borrowSockets(chunks[n], parent);
            }
        } else {
            // take the next available port
            int at = findConsecutive(1, 1);
            if (at != -1)
                out += lendItem(&items[at], parent);
        }

        return out;
    }

    void returnSockets(const QList<QUdpSocket *> &sockList)
    {
        for (QUdpSocket *sock : sockList) {
            int at = -1;
            for (int n = 0; n < items.count(); ++n) {
                if (items[n].sockList.contains(sock)) {
                    at = n;
                    break;
                }
            }

            Q_ASSERT(at != -1);

            if (at == -1)
                continue;

            Item &i = items[at];

            QHostAddress a = sock->localAddress();

            Q_ASSERT(i.lent);
            Q_ASSERT(i.lentAddrs.contains(a));

            sock->setParent(q);
            connect(sock, SIGNAL(readyRead()), SLOT(sock_readyRead()));

            i.lentAddrs.removeAll(a);
            if (i.lentAddrs.isEmpty())
                i.lent = false;
        }

        tryCleanup();
    }

private slots:
    void sock_readyRead()
    {
        QUdpSocket *sock = static_cast<QUdpSocket *>(sender());

        // eat all packets
        while (sock->hasPendingDatagrams())
            sock->readDatagram(nullptr, 0);
    }

private:
    void tryBind()
    {
        for (int n = 0; n < items.count(); ++n) {
            Item &i = items[n];

            // skip ports we don't care about
            if (!ports.contains(i.port))
                continue;

            QList<QHostAddress> neededAddrs;
            for (const QHostAddress &a : qAsConst(addrs)) {
                if (!i.haveAddress(a))
                    neededAddrs += a;
            }

            for (const QHostAddress &a : qAsConst(neededAddrs)) {
                QUdpSocket *sock = new QUdpSocket(q);

                if (!sock->bind(a, quint16(i.port))) {
                    delete sock;
                    continue;
                }

                connect(sock, SIGNAL(readyRead()), SLOT(sock_readyRead()));

                i.sockList += sock;
            }
        }
    }

    void tryCleanup()
    {
        for (int n = 0; n < items.count(); ++n) {
            Item &i = items[n];

            // don't care about this port anymore?
            if (!i.lent && !ports.contains(i.port)) {
                for (QUdpSocket *sock : qAsConst(i.sockList))
                    sock->deleteLater();

                items.removeAt(n);
                --n; // adjust position
                continue;
            }

            // any addresses we don't care about?
            for (int k = 0; k < i.sockList.count(); ++k) {
                QUdpSocket *sock = i.sockList[k];

                QHostAddress a = sock->localAddress();

                if (!addrs.contains(a) && !i.lentAddrs.contains(a)) {
                    sock->deleteLater();
                    i.sockList.removeAt(k);
                    --k; // adjust position
                    continue;
                }
            }
        }
    }

    bool isReserved(const Item &i) const
    {
        // must have desired addrs to consider a port reserved
        if (addrs.isEmpty())
            return false;

        for (const QHostAddress &a : addrs) {
            if (!i.haveAddress(a))
                return false;
        }

        return true;
    }

    bool isConsecutive(int at, int count) const
    {
        if (at + count > items.count())
            return false;

        for (int n = 0; n < count; ++n) {
            const Item &i = items[at + n];

            if (i.lent || !isReserved(i))
                return false;

            if (n > 0 && (i.port != items[at + n - 1].port + 1))
                return false;
        }

        return true;
    }

    int findConsecutive(int count, int align) const
    {
        for (int n = 0; n < items.count(); n += align) {
            if (isConsecutive(n, count))
                return n;
        }

        return -1;
    }

    QList<QUdpSocket *> lendItem(Item *i, QObject *parent)
    {
        QList<QUdpSocket *> out;

        i->lent = true;
        for (QUdpSocket *sock : qAsConst(i->sockList)) {
            i->lentAddrs += sock->localAddress();
            sock->disconnect(this);
            sock->setParent(parent);
            out += sock;
        }

        return out;
    }
};

UdpPortReserver::UdpPortReserver(QObject *parent) : QObject(parent) { d = new Private(this); }

UdpPortReserver::~UdpPortReserver() { delete d; }

void UdpPortReserver::setAddresses(const QList<QHostAddress> &addrs) { d->updateAddresses(addrs); }

void UdpPortReserver::setPorts(int start, int len)
{
    QList<int> ports;
    for (int n = 0; n < len; ++n)
        ports += start + n;
    setPorts(ports);
}

void UdpPortReserver::setPorts(const QList<int> &ports) { d->updatePorts(ports); }

bool UdpPortReserver::reservedAll() const { return d->reservedAll(); }

QList<QUdpSocket *> UdpPortReserver::borrowSockets(int portCount, QObject *parent)
{
    return d->borrowSockets(portCount, parent);
}

void UdpPortReserver::returnSockets(const QList<QUdpSocket *> &sockList) { d->returnSockets(sockList); }
} // namespace XMPP

#include "udpportreserver.moc"
