#include <gtest/gtest.h>
#include <QByteArray>
#include <QString>

#include "dns/QDnsHeader.h"

struct DNSHeaderTestParams
{
    bool isValid;
    const char* hexData;
    quint16 id;

    quint16 flags;
    bool flagQuery;
    quint8 flagOpcode;
    bool flagAuthoritative;
    bool flagTruncated;
    bool flagRecursionDesired;
    bool flagRecursionAvailable;
    bool flagReserved0;
    bool flagAnswerAuthenticated;
    bool flagNonAuthenticatedAcceptable;
    quint8 flagReplyCode;

    quint16 questionCount;
    quint16 answerCount;
    quint16 authorityCount;
    quint16 additionalCount;
};

class DNSHeaderTest : public ::testing::TestWithParam<DNSHeaderTestParams>
{
};

TEST_P(DNSHeaderTest, parseAndGetters)
{
    auto params = GetParam();
    auto data = QByteArray::fromHex(params.hexData);
    auto header = QDns::QDnsHeader::parse(data);
    ASSERT_EQ(header.isValid(), params.isValid);

    if (header.isValid())
    {
        ASSERT_EQ(header.getId(), params.id);

        ASSERT_EQ(header.getFlags(), params.flags);
        ASSERT_EQ(header.isQuery(), params.flagQuery);
        ASSERT_EQ(header.getOpcode(), params.flagOpcode);
        ASSERT_EQ(header.isAuthoritative(), params.flagAuthoritative);
        ASSERT_EQ(header.isTruncated(), params.flagTruncated);
        ASSERT_EQ(header.isRecursionDesired(), params.flagRecursionDesired);
        ASSERT_EQ(header.isRecursionAvailable(), params.flagRecursionAvailable);
        ASSERT_EQ(header.isAnswerAuthenticated(), params.flagAnswerAuthenticated);
        ASSERT_EQ(header.isNonAuthenticatedAcceptable(), params.flagNonAuthenticatedAcceptable);
        ASSERT_EQ(header.getReplyCode(), params.flagReplyCode);

        ASSERT_EQ(header.getQuestionCount(), params.questionCount);
        ASSERT_EQ(header.getAnswerCount(), params.answerCount);
        ASSERT_EQ(header.getAuthorityCount(), params.authorityCount);
        ASSERT_EQ(header.getAdditionalCount(), params.additionalCount);

        ASSERT_EQ(header.serialize().toHex(), QString(params.hexData).remove(" "));
    }
}

TEST_P(DNSHeaderTest, settersAndGetters1)
{
    auto params = GetParam();

    if (!params.isValid)
    {
        return;
    }

    auto data = QByteArray::fromHex(params.hexData);
    QDns::QDnsHeader header;

    header.setId(params.id);
    header.setFlags(params.flags);
    header.setQuestionCount(params.questionCount);
    header.setAnswerCount(params.answerCount);
    header.setAuthorityCount(params.authorityCount);
    header.setAdditionalCount(params.additionalCount);

    ASSERT_EQ(header.getId(), params.id);

    ASSERT_EQ(header.getFlags(), params.flags);
    ASSERT_EQ(header.isQuery(), params.flagQuery);
    ASSERT_EQ(header.getOpcode(), params.flagOpcode);
    ASSERT_EQ(header.isAuthoritative(), params.flagAuthoritative);
    ASSERT_EQ(header.isTruncated(), params.flagTruncated);
    ASSERT_EQ(header.isRecursionDesired(), params.flagRecursionDesired);
    ASSERT_EQ(header.isRecursionAvailable(), params.flagRecursionAvailable);
    ASSERT_EQ(header.isAnswerAuthenticated(), params.flagAnswerAuthenticated);
    ASSERT_EQ(header.isNonAuthenticatedAcceptable(), params.flagNonAuthenticatedAcceptable);
    ASSERT_EQ(header.getReplyCode(), params.flagReplyCode);

    ASSERT_EQ(header.getQuestionCount(), params.questionCount);
    ASSERT_EQ(header.getAnswerCount(), params.answerCount);
    ASSERT_EQ(header.getAuthorityCount(), params.authorityCount);
    ASSERT_EQ(header.getAdditionalCount(), params.additionalCount);

    ASSERT_EQ(header.serialize().toHex(), QString(params.hexData).remove(" "));
}

TEST_P(DNSHeaderTest, settersAndGetters2)
{
    auto params = GetParam();

    if (!params.isValid)
    {
        return;
    }

    auto data = QByteArray::fromHex(params.hexData);
    QDns::QDnsHeader header;

    header.setId(params.id);
    header.setIsQuery(params.flagQuery);
    header.setOpcode(params.flagOpcode);
    header.setIsAuthoritative(params.flagAuthoritative);
    header.setIsTruncated(params.flagTruncated);
    header.setIsRecursionDesired(params.flagRecursionDesired);
    header.setIsRecursionAvailable(params.flagRecursionAvailable);
    header.setFlagBits(9, (quint16) params.flagReserved0, 1);
    header.setIsAnswerAuthenticated(params.flagAnswerAuthenticated);
    header.setNonAuthenticatedAcceptable(params.flagNonAuthenticatedAcceptable);
    header.setReplyCode(params.flagReplyCode);
    header.setQuestionCount(params.questionCount);
    header.setAnswerCount(params.answerCount);
    header.setAuthorityCount(params.authorityCount);
    header.setAdditionalCount(params.additionalCount);

    ASSERT_EQ(header.getId(), params.id);

    ASSERT_EQ(header.getFlags(), params.flags);
    ASSERT_EQ(header.isQuery(), params.flagQuery);
    ASSERT_EQ(header.getOpcode(), params.flagOpcode);
    ASSERT_EQ(header.isAuthoritative(), params.flagAuthoritative);
    ASSERT_EQ(header.isTruncated(), params.flagTruncated);
    ASSERT_EQ(header.isRecursionDesired(), params.flagRecursionDesired);
    ASSERT_EQ(header.isRecursionAvailable(), params.flagRecursionAvailable);
    ASSERT_EQ(header.isAnswerAuthenticated(), params.flagAnswerAuthenticated);
    ASSERT_EQ(header.isNonAuthenticatedAcceptable(), params.flagNonAuthenticatedAcceptable);
    ASSERT_EQ(header.getReplyCode(), params.flagReplyCode);

    ASSERT_EQ(header.getQuestionCount(), params.questionCount);
    ASSERT_EQ(header.getAnswerCount(), params.answerCount);
    ASSERT_EQ(header.getAuthorityCount(), params.authorityCount);
    ASSERT_EQ(header.getAdditionalCount(), params.additionalCount);

    ASSERT_EQ(header.serialize().toHex(), QString(params.hexData).remove(" "));
}

INSTANTIATE_TEST_CASE_P(DNSHeaderTestInst,
                        DNSHeaderTest,
                        ::testing::Values(
                                DNSHeaderTestParams {false, ""},
                                DNSHeaderTestParams {false, "00"},
                                DNSHeaderTestParams {true,
                                                     "3c dc 01 00 00 01 00 00  00 00 00 00",
                                                     0x3cdc,

                                                     0x0100,
                                                     false,
                                                     0,
                                                     false,
                                                     false,
                                                     true,
                                                     false,
                                                     false,
                                                     false,
                                                     false,
                                                     0,

                                                     1,
                                                     0,
                                                     0,
                                                     0
                                },
                                DNSHeaderTestParams {true,
                                                     "3c dc 81 80 00 01 00 01  00 00 00 00",
                                                     0x3cdc,
                                                     0x8180,
                                                     true,
                                                     0,
                                                     false,
                                                     false,
                                                     true,
                                                     true,
                                                     false,
                                                     false,
                                                     false,
                                                     0,

                                                     1,
                                                     1,
                                                     0,
                                                     0
                                },

                                DNSHeaderTestParams {true,
                                                     "3c dc aa aa aa aa 55 55  55 ff ff 55",
                                                     0x3cdc,
                                                     0xaaaa,
                                                     true,
                                                     5,
                                                     false,
                                                     true,
                                                     false,
                                                     true,
                                                     false,
                                                     true,
                                                     false,
                                                     10,

                                                     43690,
                                                     21845,
                                                     22015,
                                                     65365
                                },
                                DNSHeaderTestParams {true,
                                                     "3c dc 55 55 55 ff aa aa  ff 55 55 55",
                                                     0x3cdc,
                                                     0x5555,
                                                     false,
                                                     10,
                                                     true,
                                                     false,
                                                     true,
                                                     false,
                                                     true,
                                                     false,
                                                     true,
                                                     5,

                                                     22015,
                                                     43690,
                                                     65365,
                                                     21845
                                }
                        ));

TEST(DNSHeaderTest, getFlagBits)
{
    QDns::QDnsHeader header;

    header.setFlags(0x5555);
    ASSERT_EQ(header.getFlagBits(0, 1), 0);
    ASSERT_EQ(header.getFlagBits(3, 4), 5);
}

TEST(DNSHeaderTest, getFlagBitIndex)
{
    ASSERT_EQ(QDns::QDnsHeader::getFlagBitIndex(0), 15);
    ASSERT_EQ(QDns::QDnsHeader::getFlagBitIndex(1), 14);
    ASSERT_EQ(QDns::QDnsHeader::getFlagBitIndex(14), 1);
    ASSERT_EQ(QDns::QDnsHeader::getFlagBitIndex(15), 0);
}

int main(int argc, char* argv[])
{
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}