#include <cstdio>
#include <cstring>
#include <vector>
#include "solution.cpp"

static int failures = 0;
#define CHECK(cond, name) do { if (cond) std::printf("ok   %s\n", name); \
    else { std::printf("FAIL %s (line %d)\n", name, __LINE__); ++failures; } } while (0)

// Независимая реализация суммы в дополнительном коде для проверки тестов.
static uint16_t ref_sum(const uint8_t* b, size_t n, uint32_t init) {
    uint32_t acc = init;
    for (size_t i = 0; i + 1 < n; i += 2) acc += (uint32_t(b[i]) << 8) | b[i + 1];
    if (n % 2) acc += uint32_t(b[n - 1]) << 8;
    while (acc >> 16) acc = (acc & 0xFFFF) + (acc >> 16);
    return uint16_t(acc);
}

// Пакет: заголовок (ihl=5) + payload_len байт нагрузки.
static std::vector<uint8_t> make_packet(size_t payload_len, uint8_t proto, bool with_options = false,
                                        bool corrupt_checksum = false, uint16_t force_total = 0) {
    size_t hlen = with_options ? 24 : 20;
    std::vector<uint8_t> p(hlen + payload_len, 0);
    p[0] = 0x40 | uint8_t(hlen / 4);           // version 4, ihl
    p[1] = 0x00;                                // DSCP/ECN
    uint16_t total = force_total ? force_total : uint16_t(hlen + payload_len);
    p[2] = uint8_t(total >> 8); p[3] = uint8_t(total & 0xFF);
    p[4] = 0x12; p[5] = 0x34;                   // id
    p[6] = 0x40; p[7] = 0x00;                   // DF
    p[8] = 64;                                  // TTL
    p[9] = proto;
    p[10] = 0; p[11] = 0;                       // checksum = 0 для расчёта
    p[12] = 0x0A; p[13] = 0x00; p[14] = 0x00; p[15] = 0x01;  // 10.0.0.1
    p[16] = 0x0A; p[17] = 0x00; p[18] = 0x00; p[19] = 0x02;  // 10.0.0.2
    if (with_options) { p[20] = 0x01; p[21] = 0x01; p[22] = 0x00; p[23] = 0x00; }
    for (size_t i = 0; i < payload_len; ++i) p[hlen + i] = uint8_t('A' + (i % 26));
    uint16_t cs = uint16_t(~ref_sum(p.data(), hlen, 0));
    p[10] = uint8_t(cs >> 8); p[11] = uint8_t(cs & 0xFF);
    if (corrupt_checksum) p[10] ^= 0xFF;
    return p;
}

int main() {
    // parse: обычный пакет
    {
        auto p = make_packet(10, 6);
        Ipv4Header h{};
        bool ok = parse_ipv4(p.data(), p.size(), &h);
        CHECK(ok, "parse() обычного пакета -> true");
        CHECK(h.version == 4, "parse: version == 4");
        CHECK(h.ihl == 5, "parse: ihl == 5");
        CHECK(h.protocol == 6, "parse: protocol == 6 (TCP)");
        CHECK(h.total_length == 30, "parse: total_length == 30");
        CHECK(h.src_ip == 0x0A000001u, "parse: src_ip == 10.0.0.1 (порядок байт как на проводе)");
        CHECK(h.dst_ip == 0x0A000002u, "parse: dst_ip == 10.0.0.2");
    }
    // parse: с опциями (ihl=6)
    {
        auto p = make_packet(0, 17, true);
        Ipv4Header h{};
        bool ok = parse_ipv4(p.data(), p.size(), &h);
        CHECK(ok && h.ihl == 6 && h.total_length == 24, "parse: ihl=6 и total_length=24 (с опциями)");
    }
    // parse: отрицательные случаи
    {
        auto p = make_packet(10, 6);
        Ipv4Header h{};
        CHECK(!parse_ipv4(p.data(), 10, &h), "parse: обрезанный буфер (10 из 30) -> false");
        CHECK(!parse_ipv4(nullptr, 0, &h), "parse: пустой буфер -> false");
        auto v6 = p; v6[0] = 0x60;
        CHECK(!parse_ipv4(v6.data(), v6.size(), &h), "parse: version 6 -> false");
        auto bad_ihl = p; bad_ihl[0] = 0x43;   // ihl=3
        CHECK(!parse_ipv4(bad_ihl.data(), bad_ihl.size(), &h), "parse: ihl < 5 -> false");
        auto too_big = make_packet(10, 6, false, false, 100); // total_length больше буфера
        CHECK(!parse_ipv4(too_big.data(), too_big.size(), &h), "parse: total_length > len -> false");
        auto small_total = make_packet(10, 6, false, false, 12); // меньше самого заголовка
        CHECK(!parse_ipv4(small_total.data(), small_total.size(), &h), "parse: total_length < ihl*4 -> false");
    }
    // checksum
    {
        auto p = make_packet(10, 6);
        auto pz = p; pz[10] = 0; pz[11] = 0;      // поле обнулено — так считают контрольную сумму
        uint16_t cs = compute_checksum(pz.data(), 20);
        CHECK(cs == uint16_t(~ref_sum(pz.data(), 20, 0)), "compute_checksum: совпадает с эталонным расчётом");
        CHECK(cs == uint16_t((p[10] << 8) | p[11]), "compute_checksum: значение равно полю в корректном пакете");
        CHECK(checksum_valid(p.data(), p.size()), "checksum_valid(корректный пакет) -> true");
        auto bad = make_packet(10, 6, false, true);
        CHECK(!checksum_valid(bad.data(), bad.size()), "checksum_valid(испорченная сумма) -> false");
        auto bad_ttl = p; bad_ttl[8] ^= 0x01;
        CHECK(!checksum_valid(bad_ttl.data(), bad_ttl.size()), "checksum_valid(испорчен TTL) -> false");
        auto bad_payload = p; bad_payload[24] ^= 0xFF;
        CHECK(checksum_valid(bad_payload.data(), bad_payload.size()),
              "checksum_valid: нагрузка в контрольную сумму заголовка не входит -> true");
        // нечётная длина
        uint8_t odd[3] = {0x01, 0x02, 0x03};
        CHECK(compute_checksum(odd, 3) == uint16_t(~ref_sum(odd, 3, 0)), "compute_checksum: нечётная длина");
    }
    std::printf(failures ? "\nFAILURES: %d\n" : "\nALL PASS\n", failures);
    return failures ? 1 : 0;
}