/*
 * Copyright (C) 2007  Justin Karneges
 *
 * 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 WITHOUT ANY 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 "qdnssd.h"

#include <QHostAddress>
#include <QtCore>
#ifdef Q_OS_WIN // for ntohl
#include <windows.h>
#else
#include <netinet/in.h>
#endif

class Command {
public:
    enum Type { Query, Browse, Resolve, Reg };

    Type type;

    QString    name;      // query, resolve, reg
    int        rtype;     // query
    QString    stype;     // browse, resolve, reg
    QString    domain;    // browse, resolve, reg
    int        port;      // reg
    QByteArray txtRecord; // reg

    int  id;
    int  dnsId;
    bool error;
    bool done; // for resolve

    Command() : error(false), done(false) { }
};

static QString nameToString(const QByteArray &in)
{
    QStringList parts;
    int         at = 0;
    while (at < in.size()) {
        int len = in[at++];
        parts += QString::fromUtf8(in.mid(at, len));
        at += len;
    }
    return parts.join(".");
}

static QString recordToDesc(const QDnsSd::Record &rec)
{
    QString desc;

    if (rec.rrtype == 1) // A
    {
        quint32 *p = (quint32 *)rec.rdata.data();
        desc       = QHostAddress(ntohl(*p)).toString();
    } else if (rec.rrtype == 28) // AAAA
    {
        desc = QHostAddress((quint8 *)rec.rdata.data()).toString();
    } else if (rec.rrtype == 12) // PTR
    {
        desc = QString("[%1]").arg(nameToString(rec.rdata));
    } else
        desc = QString("%1 bytes").arg(rec.rdata.size());

    return desc;
}

static QStringList txtRecordToStringList(const QByteArray &rdata)
{
    QList<QByteArray> txtEntries = QDnsSd::parseTxtRecord(rdata);
    if (txtEntries.isEmpty())
        return QStringList();

    QStringList out;
    foreach (const QByteArray &entry, txtEntries)
        out += QString::fromUtf8(entry);
    return out;
}

static void printIndentedTxt(const QByteArray &txtRecord)
{
    QStringList list = txtRecordToStringList(txtRecord);
    if (!list.isEmpty()) {
        foreach (const QString &s, list)
            printf("   %s\n", qPrintable(s));
    } else
        printf("   (TXT parsing error)\n");
}

class App : public QObject {
    Q_OBJECT
public:
    QList<Command> commands;
    QDnsSd *       dns;

    App()
    {
        dns = new QDnsSd(this);
        connect(dns, SIGNAL(queryResult(int, QDnsSd::QueryResult)), SLOT(dns_queryResult(int, QDnsSd::QueryResult)));
        connect(dns, SIGNAL(browseResult(int, QDnsSd::BrowseResult)),
                SLOT(dns_browseResult(int, QDnsSd::BrowseResult)));
        connect(dns, SIGNAL(resolveResult(int, QDnsSd::ResolveResult)),
                SLOT(dns_resolveResult(int, QDnsSd::ResolveResult)));
        connect(dns, SIGNAL(regResult(int, QDnsSd::RegResult)), SLOT(dns_regResult(int, QDnsSd::RegResult)));
    }

public slots:
    void start()
    {
        for (int n = 0; n < commands.count(); ++n) {
            Command &c = commands[n];

            c.id = n;
            if (c.type == Command::Query) {
                printf("%2d: Query name=[%s], type=%d ...\n", c.id, qPrintable(c.name), c.rtype);
                c.dnsId = dns->query(c.name.toUtf8(), c.rtype);
            } else if (c.type == Command::Browse) {
                printf("%2d: Browse type=[%s]", c.id, qPrintable(c.stype));
                if (!c.domain.isEmpty())
                    printf(", domain=[%s]", qPrintable(c.domain));
                printf(" ...\n");
                c.dnsId = dns->browse(c.stype.toUtf8(), c.domain.toUtf8());
            } else if (c.type == Command::Resolve) {
                printf("%2d: Resolve name=[%s], type=[%s], domain=[%s] ...\n", c.id, qPrintable(c.name),
                       qPrintable(c.stype), qPrintable(c.domain));
                c.dnsId = dns->resolve(c.name.toUtf8(), c.stype.toUtf8(), c.domain.toUtf8());
            } else if (c.type == Command::Reg) {
                printf("%2d: Register name=[%s], type=[%s]", c.id, qPrintable(c.name), qPrintable(c.stype));
                if (!c.domain.isEmpty())
                    printf(", domain=[%s]", qPrintable(c.domain));
                printf(", port=%d ...\n", c.port);
                if (!c.txtRecord.isEmpty())
                    printIndentedTxt(c.txtRecord);

                c.dnsId = dns->reg(c.name.toUtf8(), c.stype.toUtf8(), c.domain.toUtf8(), c.port, c.txtRecord);
            }
        }
    }

signals:
    void quit();

private:
    int cmdIdToCmdIndex(int cmdId)
    {
        for (int n = 0; n < commands.count(); ++n) {
            const Command &c = commands[n];
            if (c.id == cmdId)
                return n;
        }
        return -1;
    }

    int dnsIdToCmdIndex(int dnsId)
    {
        for (int n = 0; n < commands.count(); ++n) {
            const Command &c = commands[n];
            if (c.dnsId == dnsId)
                return n;
        }
        return -1;
    }

    void tryQuit()
    {
        // quit if there are nothing but errors or completed resolves
        bool doQuit = true;
        foreach (const Command &c, commands) {
            if (c.error || (c.type == Command::Resolve && c.done))
                continue;

            doQuit = false;
            break;
        }

        if (doQuit)
            emit quit();
    }

private slots:
    void dns_queryResult(int id, const QDnsSd::QueryResult &result)
    {
        int      at = dnsIdToCmdIndex(id);
        Command &c  = commands[at];

        if (!result.success) {
            printf("%2d: Error.", c.id);
            if (!result.lowLevelError.func.isEmpty())
                printf(" (%s, %d)", qPrintable(result.lowLevelError.func), result.lowLevelError.code);
            printf("\n");
            c.error = true;
            tryQuit();
            return;
        }

        foreach (const QDnsSd::Record &rec, result.records) {
            if (rec.added) {
                printf("%2d: Added:   %s, ttl=%u\n", c.id, qPrintable(recordToDesc(rec)), rec.ttl);
                if (rec.rrtype == 16)
                    printIndentedTxt(rec.rdata);
            } else
                printf("%2d: Removed: %s, ttl=%u\n", c.id, qPrintable(recordToDesc(rec)), rec.ttl);
        }
    }

    void dns_browseResult(int id, const QDnsSd::BrowseResult &result)
    {
        int      at = dnsIdToCmdIndex(id);
        Command &c  = commands[at];

        if (!result.success) {
            printf("%2d: Error.", c.id);
            if (!result.lowLevelError.func.isEmpty())
                printf(" (%s, %d)", qPrintable(result.lowLevelError.func), result.lowLevelError.code);
            printf("\n");
            c.error = true;
            tryQuit();
            return;
        }

        foreach (const QDnsSd::BrowseEntry &e, result.entries) {
            if (e.added)
                printf("%2d: Added:   [%s] [%s] [%s]\n", c.id, qPrintable(QString::fromUtf8(e.serviceName)),
                       qPrintable(QString::fromUtf8(e.serviceType)), qPrintable(QString::fromUtf8(e.replyDomain)));
            else
                printf("%2d: Removed: [%s]\n", c.id, qPrintable(QString::fromUtf8(e.serviceName)));
        }
    }

    void dns_resolveResult(int id, const QDnsSd::ResolveResult &result)
    {
        int      at = dnsIdToCmdIndex(id);
        Command &c  = commands[at];

        if (!result.success) {
            printf("%2d: Error.", c.id);
            if (!result.lowLevelError.func.isEmpty())
                printf(" (%s, %d)", qPrintable(result.lowLevelError.func), result.lowLevelError.code);
            printf("\n");
            c.error = true;
            tryQuit();
            return;
        }

        printf("%2d: Result: host=[%s] port=%d\n", c.id, qPrintable(QString::fromUtf8(result.hostTarget)), result.port);
        if (!result.txtRecord.isEmpty())
            printIndentedTxt(result.txtRecord);

        c.done = true;
        tryQuit();
    }

    void dns_regResult(int id, const QDnsSd::RegResult &result)
    {
        int      at = dnsIdToCmdIndex(id);
        Command &c  = commands[at];

        if (!result.success) {
            QString errstr;
            if (result.errorCode == QDnsSd::RegResult::ErrorConflict)
                errstr = "Conflict";
            else
                errstr = "Generic";
            printf("%2d: Error (%s).", c.id, qPrintable(errstr));
            if (!result.lowLevelError.func.isEmpty())
                printf(" (%s, %d)", qPrintable(result.lowLevelError.func), result.lowLevelError.code);
            printf("\n");
            c.error = true;
            tryQuit();
            return;
        }

        printf("%2d: Registered: domain=[%s]\n", c.id, qPrintable(QString::fromUtf8(result.domain)));
    }
};

#include "sdtest.moc"

void usage()
{
    printf("usage: sdtest [[command] (command) ...]\n");
    printf(" options: --txt=str0,...,strn\n");
    printf("\n");
    printf(" q=name,type#                   query for a record\n");
    printf(" b=type(,domain)                browse for services\n");
    printf(" r=name,type,domain             resolve a service\n");
    printf(" e=name,type,port(,domain)      register a service\n");
    printf("\n");
}

int main(int argc, char **argv)
{
    QCoreApplication qapp(argc, argv);

    QStringList args = qapp.arguments();
    args.removeFirst();

    if (args.count() < 1) {
        usage();
        return 1;
    }

    // options
    QStringList txt;

    for (int n = 0; n < args.count(); ++n) {
        QString s = args[n];
        if (!s.startsWith("--"))
            continue;
        QString var;
        QString val;
        int     x = s.indexOf('=');
        if (x != -1) {
            var = s.mid(2, x - 2);
            val = s.mid(x + 1);
        } else {
            var = s.mid(2);
        }

        bool known = true;

        if (var == "txt") {
            txt = val.split(',');
        } else
            known = false;

        if (known) {
            args.removeAt(n);
            --n; // adjust position
        }
    }

    // commands
    QList<Command> commands;

    for (int n = 0; n < args.count(); ++n) {
        QString str = args[n];
        int     n   = str.indexOf('=');
        if (n == -1) {
            printf("Error: bad format of command.\n");
            return 1;
        }

        QString     type  = str.mid(0, n);
        QString     rest  = str.mid(n + 1);
        QStringList parts = rest.split(',');

        if (type == "q") {
            if (parts.count() < 2) {
                usage();
                return 1;
            }

            Command c;
            c.type  = Command::Query;
            c.name  = parts[0];
            c.rtype = parts[1].toInt();
            commands += c;
        } else if (type == "b") {
            if (parts.count() < 1) {
                usage();
                return 1;
            }

            Command c;
            c.type  = Command::Browse;
            c.stype = parts[0];
            if (parts.count() >= 2)
                c.domain = parts[1];
            commands += c;
        } else if (type == "r") {
            if (parts.count() < 3) {
                usage();
                return 1;
            }

            Command c;
            c.type   = Command::Resolve;
            c.name   = parts[0];
            c.stype  = parts[1];
            c.domain = parts[2];
            commands += c;
        } else if (type == "e") {
            if (parts.count() < 3) {
                usage();
                return 1;
            }

            Command c;
            c.type  = Command::Reg;
            c.name  = parts[0];
            c.stype = parts[1];
            c.port  = parts[2].toInt();
            if (parts.count() >= 4)
                c.domain = parts[3];

            if (!txt.isEmpty()) {
                QList<QByteArray> strings;
                foreach (const QString &str, txt)
                    strings += str.toUtf8();

                QByteArray txtRecord = QDnsSd::createTxtRecord(strings);
                if (txtRecord.isEmpty()) {
                    printf("Error: failed to create TXT record, input too large or invalid\n");
                    return 1;
                }

                c.txtRecord = txtRecord;
            }

            commands += c;
        } else {
            printf("Error: unknown command type '%s'.\n", qPrintable(type));
            return 1;
        }
    }

    App app;
    app.commands = commands;
    QObject::connect(&app, SIGNAL(quit()), &qapp, SLOT(quit()));
    QTimer::singleShot(0, &app, SLOT(start()));
    qapp.exec();
    return 0;
}
