/*
 * 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, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301  USA
 *
 */

#include <QtCore>
#include <QHostAddress>
#include "qdnssd.h"

// for ntohl
#ifdef Q_OS_WIN
# 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;
}
