#include <doctest/doctest.h>

#include <vector>

#include "packeteer/l7/dns.hpp"

using namespace packeteer::net;

namespace {

// "example.com" A query, id=0x129d - the same shape as the real query
// captured live over tailscale0 while testing the DNS dissector against
// tshark (id 0x129d / 4765 matched tshark's independent decode exactly).
std::vector<unsigned char> example_com_query() {
    return {
        0x12, 0x9d,              // id = 4765
        0x01, 0x00,              // flags: RD=1
        0x00, 0x01,              // qdcount = 1
        0x00, 0x00,              // ancount = 0
        0x00, 0x00,              // nscount = 0
        0x00, 0x00,              // arcount = 0
        7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0,
        0x00, 0x01,  // qtype = A
        0x00, 0x01,  // qclass = IN
    };
}

void append_be16(std::vector<unsigned char>& out, std::uint16_t v) {
    out.push_back(static_cast<unsigned char>(v >> 8));
    out.push_back(static_cast<unsigned char>(v & 0xFF));
}

void append_be32(std::vector<unsigned char>& out, std::uint32_t v) {
    out.push_back(static_cast<unsigned char>(v >> 24));
    out.push_back(static_cast<unsigned char>(v >> 16));
    out.push_back(static_cast<unsigned char>(v >> 8));
    out.push_back(static_cast<unsigned char>(v & 0xFF));
}

// Builds a real, well-formed DNS response for "example.com" A, with
// `answers` real resource records appended after the question - each
// using a compressed name pointer back to the question's name (offset
// 12, right after the header), exactly how real DNS servers answer,
// rather than repeating the literal name.
struct AnswerSpec {
    std::uint16_t type;
    std::uint32_t ttl;
    std::vector<unsigned char> rdata;
};

std::vector<unsigned char> build_dns_response(const std::vector<AnswerSpec>& answers) {
    std::vector<unsigned char> bytes = {
        0x12, 0x9d,
        0x81, 0x80,  // flags: QR=1 (response), RD=1, RA=1
    };
    append_be16(bytes, 1);  // qdcount
    append_be16(bytes, static_cast<std::uint16_t>(answers.size()));
    append_be16(bytes, 0);  // nscount
    append_be16(bytes, 0);  // arcount

    bytes.insert(bytes.end(), {7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0});
    append_be16(bytes, 1);  // qtype A
    append_be16(bytes, 1);  // qclass IN

    for (const auto& answer : answers) {
        bytes.push_back(0xC0);
        bytes.push_back(0x0C);  // NAME: pointer to offset 12 (the question's name)
        append_be16(bytes, answer.type);
        append_be16(bytes, 1);  // CLASS: IN
        append_be32(bytes, answer.ttl);
        append_be16(bytes, static_cast<std::uint16_t>(answer.rdata.size()));
        bytes.insert(bytes.end(), answer.rdata.begin(), answer.rdata.end());
    }
    return bytes;
}

}  // namespace

TEST_CASE("parse_dns decodes a query") {
    auto msg = parse_dns(example_com_query());
    REQUIRE(msg.has_value());
    CHECK(msg->header.id == 4765);
    CHECK_FALSE(msg->header.is_response);
    CHECK(msg->header.qdcount == 1);
    REQUIRE(msg->question.has_value());
    CHECK(msg->question->name == "example.com");
    CHECK(msg->question->qtype == 1);
}

TEST_CASE("parse_dns decodes a response") {
    std::vector<unsigned char> bytes = {
        0x12, 0x9d,
        0x81, 0x80,  // flags: QR=1 (response), RD=1, RA=1
        0x00, 0x01,  // qdcount = 1
        0x00, 0x02,  // ancount = 2
        0x00, 0x00,
        0x00, 0x00,
        7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0,
        0x00, 0x01, 0x00, 0x01,
    };
    auto msg = parse_dns(bytes);
    REQUIRE(msg.has_value());
    CHECK(msg->header.is_response);
    CHECK(msg->header.ancount == 2);
}

TEST_CASE("parse_dns rejects a truncated header") {
    std::vector<unsigned char> bytes(5, 0);
    CHECK_FALSE(parse_dns(bytes).has_value());
}

TEST_CASE("read_dns_name rejects a compression pointer") {
    std::vector<unsigned char> bytes = {0xC0, 0x0C};  // pointer: unsupported by design
    CHECK_FALSE(read_dns_name(bytes, 0).has_value());
}

TEST_CASE("DnsDissector claims port 53 and its summary matches parse_dns") {
    DnsDissector dissector;
    CHECK(dissector.port() == kDnsPort);

    auto summary = dissector.summarize(example_com_query());
    REQUIRE(summary.has_value());
    CHECK(summary->substr(0, 9) == "DNS query");
    CHECK(summary->find("example.com") != std::string::npos);
}

TEST_CASE("DnsDissector::summarize returns nullopt for a truncated payload") {
    DnsDissector dissector;
    std::vector<unsigned char> bytes(5, 0);
    CHECK_FALSE(dissector.summarize(bytes).has_value());
}

TEST_CASE("read_dns_name_following_pointers follows a compressed name back to the question") {
    auto bytes = build_dns_response({{kDnsTypeA, 300, {93, 184, 216, 34}}});
    // The answer's NAME field is the 2-byte pointer right after the
    // question section (byte offset 12 + 17 = 29 in this layout).
    auto result = read_dns_name_following_pointers(bytes, 29);
    REQUIRE(result.has_value());
    CHECK(result->first == "example.com");
}

TEST_CASE("read_dns_name_following_pointers is bounded against a pointer cycle") {
    // Two pointers pointing at each other - a backward-only check
    // wouldn't catch this (pointer B points backward to A, which
    // points forward to B), but the jump-count bound does.
    std::vector<unsigned char> bytes = {0xC0, 0x02, 0xC0, 0x00};
    CHECK_FALSE(read_dns_name_following_pointers(bytes, 0).has_value());
}

TEST_CASE("parse_dns decodes a single A answer") {
    auto bytes = build_dns_response({{kDnsTypeA, 300, {93, 184, 216, 34}}});
    auto msg = parse_dns(bytes);
    REQUIRE(msg.has_value());
    REQUIRE(msg->answers.size() == 1);
    CHECK(msg->answers[0].name == "example.com");
    CHECK(msg->answers[0].type == kDnsTypeA);
    CHECK(msg->answers[0].ttl == 300);
    REQUIRE(msg->answers[0].rdata_text.has_value());
    CHECK(*msg->answers[0].rdata_text == "93.184.216.34");
}

TEST_CASE("parse_dns decodes multiple answers, e.g. a CNAME followed by an A record") {
    std::vector<unsigned char> cname_rdata = {3, 'w', 'w', 'w', 0xC0, 0x0C};  // "www" + pointer
    auto bytes = build_dns_response(
        {{kDnsTypeCname, 60, cname_rdata}, {kDnsTypeA, 300, {93, 184, 216, 34}}});
    auto msg = parse_dns(bytes);
    REQUIRE(msg.has_value());
    REQUIRE(msg->answers.size() == 2);
    REQUIRE(msg->answers[0].rdata_text.has_value());
    CHECK(*msg->answers[0].rdata_text == "www.example.com");
    REQUIRE(msg->answers[1].rdata_text.has_value());
    CHECK(*msg->answers[1].rdata_text == "93.184.216.34");
}

TEST_CASE("parse_dns decodes an AAAA answer") {
    std::vector<unsigned char> aaaa_rdata = {0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0,
                                               0,    0,    0,    0,    0, 0, 0, 1};
    auto bytes = build_dns_response({{kDnsTypeAaaa, 300, aaaa_rdata}});
    auto msg = parse_dns(bytes);
    REQUIRE(msg.has_value());
    REQUIRE(msg->answers.size() == 1);
    REQUIRE(msg->answers[0].rdata_text.has_value());
    CHECK(*msg->answers[0].rdata_text == "2001:0db8:0000:0000:0000:0000:0000:0001");
}

TEST_CASE("parse_dns leaves rdata_text unset for an undecoded record type") {
    auto bytes = build_dns_response({{15 /* MX */, 60, {0, 10, 4, 'm', 'a', 'i', 'l'}}});
    auto msg = parse_dns(bytes);
    REQUIRE(msg.has_value());
    REQUIRE(msg->answers.size() == 1);
    CHECK(msg->answers[0].type == 15);
    CHECK_FALSE(msg->answers[0].rdata_text.has_value());
}

TEST_CASE("parse_dns stops decoding answers on the first malformed record") {
    auto bytes = build_dns_response({{kDnsTypeA, 300, {93, 184, 216, 34}}});
    bytes.resize(bytes.size() - 2);  // truncate the last answer's rdata
    auto msg = parse_dns(bytes);
    REQUIRE(msg.has_value());
    CHECK(msg->answers.empty());
    CHECK(msg->header.ancount == 1);  // the header claim is preserved even though decode failed
}

TEST_CASE("DnsDissector::summarize includes resolved addresses for a response") {
    DnsDissector dissector;
    auto bytes = build_dns_response({{kDnsTypeA, 300, {93, 184, 216, 34}}});
    auto summary = dissector.summarize(bytes);
    REQUIRE(summary.has_value());
    CHECK(summary->find("-> 93.184.216.34") != std::string::npos);
}