#include <doctest/doctest.h>

#include <vector>

#include "packeteer/net/checksum.hpp"

using namespace packeteer::net;

namespace {

// Mirrors checksum.hpp's own detail::build_ipv4_pseudo_header, kept
// separate here deliberately: constructing expected test vectors using
// the exact same private helper the code under test uses would make
// these tests circular. A few duplicated lines of test-only setup is
// the honest cost of testing independently.
std::vector<unsigned char> pseudo_header(const std::array<unsigned char, 4>& src,
                                          const std::array<unsigned char, 4>& dst,
                                          unsigned char protocol,
                                          std::span<const unsigned char> segment) {
    std::vector<unsigned char> buf;
    buf.insert(buf.end(), src.begin(), src.end());
    buf.insert(buf.end(), dst.begin(), dst.end());
    buf.push_back(0);
    buf.push_back(protocol);
    std::uint16_t len = static_cast<std::uint16_t>(segment.size());
    buf.push_back(static_cast<unsigned char>(len >> 8));
    buf.push_back(static_cast<unsigned char>(len & 0xFF));
    buf.insert(buf.end(), segment.begin(), segment.end());
    return buf;
}

}  // namespace

TEST_CASE("internet_checksum matches RFC 1071's own worked example") {
    // The RFC's example data (0001 f203 f4f5 f6f7) computes to checksum
    // 220d - an external reference, not derived from this code.
    std::vector<unsigned char> data = {0x00, 0x01, 0xf2, 0x03, 0xf4, 0xf5, 0xf6, 0xf7};
    CHECK(internet_checksum(data) == 0x220d);
}

TEST_CASE("internet_checksum of data with its own valid checksum appended is zero") {
    // Direct consequence of the RFC 1071 example: appending that
    // checksum as one more word should sum to all-ones, complementing
    // to exactly zero - this is the actual verification technique
    // verify_ipv4_checksum() etc. rely on.
    std::vector<unsigned char> data = {0x00, 0x01, 0xf2, 0x03, 0xf4, 0xf5, 0xf6, 0xf7, 0x22, 0x0d};
    CHECK(internet_checksum(data) == 0);
}

TEST_CASE("internet_checksum handles an odd-length buffer (trailing byte padded high)") {
    std::vector<unsigned char> data = {0x00, 0x01, 0xf2};  // 3 bytes: one word + one odd byte
    // 0x0001 + 0xf200 (odd byte in the high half) = 0xf201; ~0xf201 = 0x0dfe
    CHECK(internet_checksum(data) == 0x0dfe);
}

TEST_CASE("verify_ipv4_checksum accepts a header with a correctly computed checksum") {
    std::vector<unsigned char> header(20, 0);
    header[0] = 0x45;
    header[8] = 64;  // ttl
    header[9] = kProtoTcp;
    header[12] = 10; header[13] = 0; header[14] = 0; header[15] = 1;
    header[16] = 10; header[17] = 0; header[18] = 0; header[19] = 2;
    // checksum field (bytes 10-11) computed with itself still zeroed
    std::uint16_t csum = internet_checksum(header);
    header[10] = static_cast<unsigned char>(csum >> 8);
    header[11] = static_cast<unsigned char>(csum & 0xFF);

    CHECK(verify_ipv4_checksum(header));
}

TEST_CASE("verify_ipv4_checksum rejects a header corrupted after the checksum was computed") {
    std::vector<unsigned char> header(20, 0);
    header[0] = 0x45;
    header[9] = kProtoTcp;
    std::uint16_t csum = internet_checksum(header);
    header[10] = static_cast<unsigned char>(csum >> 8);
    header[11] = static_cast<unsigned char>(csum & 0xFF);

    header[15] ^= 0xFF;  // flip a source-address byte after the fact
    CHECK_FALSE(verify_ipv4_checksum(header));
}

TEST_CASE("verify_tcp_checksum_ipv4 accepts a segment with a correctly computed checksum") {
    std::array<unsigned char, 4> src = {10, 0, 0, 1};
    std::array<unsigned char, 4> dst = {10, 0, 0, 2};

    std::vector<unsigned char> tcp(20, 0);
    tcp[0] = 0; tcp[1] = 80;      // src port
    tcp[2] = 0x01; tcp[3] = 0xbb;  // dst port 443
    tcp[12] = 5 << 4;             // data_offset = 5

    auto buf = pseudo_header(src, dst, kProtoTcp, tcp);
    std::uint16_t csum = internet_checksum(buf);
    tcp[16] = static_cast<unsigned char>(csum >> 8);
    tcp[17] = static_cast<unsigned char>(csum & 0xFF);

    CHECK(verify_tcp_checksum_ipv4({src}, {dst}, tcp) == ChecksumResult::kValid);
}

TEST_CASE("verify_tcp_checksum_ipv4 rejects a segment corrupted after the checksum was computed") {
    std::array<unsigned char, 4> src = {10, 0, 0, 1};
    std::array<unsigned char, 4> dst = {10, 0, 0, 2};

    std::vector<unsigned char> tcp(20, 0);
    tcp[12] = 5 << 4;
    auto buf = pseudo_header(src, dst, kProtoTcp, tcp);
    std::uint16_t csum = internet_checksum(buf);
    tcp[16] = static_cast<unsigned char>(csum >> 8);
    tcp[17] = static_cast<unsigned char>(csum & 0xFF);

    tcp[0] ^= 0xFF;  // corrupt the source port after the fact
    CHECK(verify_tcp_checksum_ipv4({src}, {dst}, tcp) == ChecksumResult::kInvalid);
}

TEST_CASE("verify_udp_checksum_ipv4 treats a transmitted checksum of 0x0000 as not present") {
    std::array<unsigned char, 4> src = {10, 0, 0, 1};
    std::array<unsigned char, 4> dst = {10, 0, 0, 2};
    std::vector<unsigned char> udp = {0x00, 0x35, 0x00, 0x35, 0x00, 0x08, 0x00, 0x00};  // csum=0
    CHECK(verify_udp_checksum_ipv4({src}, {dst}, udp) == ChecksumResult::kNotPresent);
}

TEST_CASE("verify_udp_checksum_ipv4 accepts a datagram with a correctly computed checksum") {
    std::array<unsigned char, 4> src = {10, 0, 0, 1};
    std::array<unsigned char, 4> dst = {10, 0, 0, 2};

    std::vector<unsigned char> udp = {0x00, 0x35, 0x00, 0x35, 0x00, 0x08, 0x00, 0x00};
    auto buf = pseudo_header(src, dst, kProtoUdp, udp);
    std::uint16_t csum = internet_checksum(buf);
    // A computed checksum of exactly 0 is itself sent as 0xFFFF per
    // RFC 768, to keep it distinguishable from "no checksum" - not
    // exercised by this test's specific values, but worth the note.
    udp[6] = static_cast<unsigned char>(csum >> 8);
    udp[7] = static_cast<unsigned char>(csum & 0xFF);

    CHECK(verify_udp_checksum_ipv4({src}, {dst}, udp) == ChecksumResult::kValid);
}