#include <doctest/doctest.h>

#include <vector>

#include "packeteer/net/ipv4.hpp"  // for kProtoTcp
#include "packeteer/net/ipv6.hpp"

using namespace packeteer::net;

namespace {

Ipv6Address addr_from_groups(std::array<std::uint16_t, 8> groups) {
    Ipv6Address addr{};
    for (std::size_t i = 0; i < 8; ++i) {
        addr.bytes[i * 2] = static_cast<unsigned char>(groups[i] >> 8);
        addr.bytes[i * 2 + 1] = static_cast<unsigned char>(groups[i] & 0xFF);
    }
    return addr;
}

}  // namespace

TEST_CASE("parse_ipv6 decodes header fields and leaves the right payload") {
    std::vector<unsigned char> bytes(40, 0);
    bytes[0] = 0x60;  // version 6, traffic class high nibble 0
    bytes[1] = 0x00;  // traffic class low nibble 0, flow label starts 0
    bytes[4] = 0x00;
    bytes[5] = 0x04;  // payload_length = 4
    bytes[6] = kProtoTcp;
    bytes[7] = 64;  // hop_limit
    // src = 2001:0db8::1
    bytes[8] = 0x20; bytes[9] = 0x01; bytes[10] = 0x0d; bytes[11] = 0xb8;
    bytes[23] = 0x01;
    // dst = ::1
    bytes[39] = 0x01;
    bytes.insert(bytes.end(), {0xAA, 0xBB, 0xCC, 0xDD});

    auto ip6 = parse_ipv6(bytes);
    REQUIRE(ip6.has_value());
    CHECK(ip6->header.version == 6);
    CHECK(ip6->header.payload_length == 4);
    CHECK(ip6->header.next_header == kProtoTcp);
    CHECK(ip6->header.hop_limit == 64);
    REQUIRE(ip6->payload.size() == 4);
    CHECK(ip6->payload[0] == 0xAA);
}

TEST_CASE("parse_ipv6 rejects a non-IPv6 version") {
    std::vector<unsigned char> bytes(40, 0);
    bytes[0] = 0x45;  // version 4
    CHECK_FALSE(parse_ipv6(bytes).has_value());
}

TEST_CASE("parse_ipv6 rejects a buffer shorter than the 40-byte header") {
    std::vector<unsigned char> bytes(39, 0);
    bytes[0] = 0x60;
    CHECK_FALSE(parse_ipv6(bytes).has_value());
}

TEST_CASE("ipv6_to_string compresses the loopback address") {
    CHECK(ipv6_to_string(addr_from_groups({0, 0, 0, 0, 0, 0, 0, 1})) == "::1");
}

TEST_CASE("ipv6_to_string compresses the unspecified address") {
    CHECK(ipv6_to_string(addr_from_groups({0, 0, 0, 0, 0, 0, 0, 0})) == "::");
}

TEST_CASE("ipv6_to_string compresses a zero run in the middle") {
    CHECK(ipv6_to_string(addr_from_groups({0x2001, 0x0db8, 0, 0, 0, 0, 0, 1})) == "2001:db8::1");
}

TEST_CASE("ipv6_to_string does not compress a lone zero group") {
    CHECK(ipv6_to_string(addr_from_groups({0x2001, 0, 0x0db8, 1, 1, 1, 1, 1})) ==
          "2001:0:db8:1:1:1:1:1");
}

TEST_CASE("ipv6_to_string picks the leftmost run when two runs tie in length") {
    // Two runs of length 2: groups[1..2] and groups[5..6]. Leftmost wins.
    CHECK(ipv6_to_string(addr_from_groups({1, 0, 0, 2, 3, 0, 0, 4})) == "1::2:3:0:0:4");
}

TEST_CASE("ipv6_to_string leaves an address with no zero run untouched") {
    CHECK(ipv6_to_string(addr_from_groups({1, 2, 3, 4, 5, 6, 7, 8})) == "1:2:3:4:5:6:7:8");
}

TEST_CASE("walk_ipv6_extension_headers passes a direct transport protocol through unchanged") {
    std::vector<unsigned char> payload = {0xAA, 0xBB, 0xCC};
    auto result = walk_ipv6_extension_headers(kProtoTcp, payload);
    CHECK(result.final_next_header == kProtoTcp);
    CHECK_FALSE(result.stopped_at_esp);
    REQUIRE(result.payload.size() == 3);
    CHECK(result.payload[0] == 0xAA);
}

TEST_CASE("walk_ipv6_extension_headers walks a single Hop-by-Hop header to reach TCP") {
    // Hop-by-Hop: next_header(1)=TCP, hdr_ext_len(1)=0 -> total len (0+1)*8=8 bytes.
    std::vector<unsigned char> payload = {static_cast<unsigned char>(kProtoTcp), 0x00,
                                           0, 0, 0, 0, 0, 0};  // 6 bytes of option padding
    std::vector<unsigned char> tcp_marker = {0xDE, 0xAD};
    payload.insert(payload.end(), tcp_marker.begin(), tcp_marker.end());

    auto result = walk_ipv6_extension_headers(kNextHeaderHopByHop, payload);
    CHECK(result.final_next_header == kProtoTcp);
    CHECK_FALSE(result.stopped_at_esp);
    REQUIRE(result.payload.size() == 2);
    CHECK(result.payload[0] == 0xDE);
}

TEST_CASE("walk_ipv6_extension_headers walks a chain of two extension headers") {
    // Hop-by-Hop (8 bytes) -> Destination Options (8 bytes) -> UDP.
    std::vector<unsigned char> payload = {
        kNextHeaderDestOptions, 0x00, 0, 0, 0, 0, 0, 0,          // Hop-by-Hop, len 8
        static_cast<unsigned char>(kProtoUdp), 0x00, 0, 0, 0, 0, 0, 0,  // Dest Options, len 8
        0xFE, 0xED,                                               // "UDP header" marker
    };
    auto result = walk_ipv6_extension_headers(kNextHeaderHopByHop, payload);
    CHECK(result.final_next_header == kProtoUdp);
    REQUIRE(result.payload.size() == 2);
    CHECK(result.payload[0] == 0xFE);
}

TEST_CASE("walk_ipv6_extension_headers walks the fixed-size Fragment header") {
    std::vector<unsigned char> payload = {static_cast<unsigned char>(kProtoTcp), 0x00,
                                           0x00, 0x00, 0x00, 0x00, 0x00, 0x01,  // 8-byte fragment header
                                           0xCA, 0xFE};
    auto result = walk_ipv6_extension_headers(kNextHeaderFragment, payload);
    CHECK(result.final_next_header == kProtoTcp);
    REQUIRE(result.payload.size() == 2);
    CHECK(result.payload[0] == 0xCA);
}

TEST_CASE("walk_ipv6_extension_headers reports the fragment id even for the first fragment") {
    // Same shape as the test above (fragment offset 0 - the first
    // fragment), but this time checking that the walk continues on to
    // TCP *and* still surfaces the Identification field, which a
    // caller needs to correlate this with the fragments that follow.
    std::vector<unsigned char> payload = {static_cast<unsigned char>(kProtoTcp), 0x00,
                                           0x00, 0x00, 0x00, 0x00, 0x30, 0x39,  // id = 0x3039
                                           0xCA, 0xFE};
    auto result = walk_ipv6_extension_headers(kNextHeaderFragment, payload);
    CHECK_FALSE(result.is_non_first_fragment);
    REQUIRE(result.fragment_id.has_value());
    CHECK(*result.fragment_id == 0x3039);
    CHECK(result.final_next_header == kProtoTcp);
    REQUIRE(result.payload.size() == 2);  // walk continued past the fragment header to real payload
}

TEST_CASE("walk_ipv6_extension_headers stops at a non-first fragment rather than walking into "
          "continuation data") {
    // Fragment offset field (13 bits, packed into the top of bytes[2:3])
    // set to a nonzero value - 8 in units of 8 bytes, i.e. byte offset
    // 64 into the original datagram. offset_res_m = 8 << 3 = 0x0040.
    std::vector<unsigned char> payload = {
        static_cast<unsigned char>(kProtoTcp), 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x2A,
        0xDE, 0xAD, 0xBE, 0xEF,  // pure continuation data - NOT a TCP header
    };
    auto result = walk_ipv6_extension_headers(kNextHeaderFragment, payload);
    CHECK(result.is_non_first_fragment);
    REQUIRE(result.fragment_id.has_value());
    CHECK(*result.fragment_id == 0x2A);
    // final_next_header still names TCP (that's what the reassembled
    // datagram eventually is), but the payload past it is untouched
    // continuation data - callers must not decode it as TCP.
    CHECK(result.final_next_header == kProtoTcp);
    REQUIRE(result.payload.size() == 4);
    CHECK(result.payload[0] == 0xDE);
}

TEST_CASE("walk_ipv6_extension_headers applies AH's 4-byte-unit length formula") {
    // AH: next_header(1)=TCP, payload_len(1)=1 -> total len (1+2)*4=12 bytes.
    std::vector<unsigned char> payload(12, 0);
    payload[0] = static_cast<unsigned char>(kProtoTcp);
    payload[1] = 0x01;
    payload.push_back(0x11);
    payload.push_back(0x22);

    auto result = walk_ipv6_extension_headers(kNextHeaderAh, payload);
    CHECK(result.final_next_header == kProtoTcp);
    REQUIRE(result.payload.size() == 2);
    CHECK(result.payload[0] == 0x11);
}

TEST_CASE("walk_ipv6_extension_headers stops at ESP without guessing past it") {
    std::vector<unsigned char> payload = {0x01, 0x02, 0x03, 0x04};
    auto result = walk_ipv6_extension_headers(kNextHeaderEsp, payload);
    CHECK(result.stopped_at_esp);
    CHECK(result.final_next_header == kNextHeaderEsp);
    REQUIRE(result.payload.size() == 4);
    CHECK(result.payload[0] == 0x01);  // untouched: ESP payload starts right here
}

TEST_CASE("walk_ipv6_extension_headers stops gracefully on a truncated extension header") {
    std::vector<unsigned char> payload = {static_cast<unsigned char>(kProtoTcp),
                                           0xFF};  // claims (255+1)*8 bytes; nowhere near present
    auto result = walk_ipv6_extension_headers(kNextHeaderHopByHop, payload);
    CHECK(result.final_next_header == kNextHeaderHopByHop);  // never resolved past it
    CHECK_FALSE(result.stopped_at_esp);
}

TEST_CASE("walk_ipv6_extension_headers passes an unknown next_header through untouched") {
    std::vector<unsigned char> payload = {0x01, 0x02};
    auto result = walk_ipv6_extension_headers(200, payload);  // not a known extension type
    CHECK(result.final_next_header == 200);
    REQUIRE(result.payload.size() == 2);
    CHECK(result.payload[0] == 0x01);
}