#include <doctest/doctest.h>

#include <cstdio>
#include <vector>

#include "wireframe/pcapng/reader.hpp"
#include "wireframe/pcapng/writer.hpp"

using namespace wireframe::pcapng;

TEST_CASE("pcapng writer/reader round-trip a single packet") {
    std::FILE* f = std::tmpfile();
    REQUIRE(f != nullptr);

    Writer writer(f);
    writer.write_section_header();
    writer.write_interface_description(65535, kLinkTypeEthernet);

    std::vector<unsigned char> packet_data = {0xDE, 0xAD, 0xBE, 0xEF, 0x00};
    writer.write_packet(/*interface_id=*/0, /*ts_sec=*/1700000000, /*ts_usec=*/123456,
                         packet_data, /*original_len=*/5);

    std::fflush(f);
    std::fseek(f, 0, SEEK_SET);

    Reader reader(f);
    auto record = reader.next_packet();
    REQUIRE(record.has_value());
    CHECK(record->interface_id == 0);
    CHECK(record->timestamp_us == 1700000000ULL * 1'000'000ULL + 123456ULL);
    CHECK(record->original_len == 5);
    CHECK(record->data == packet_data);

    CHECK_FALSE(reader.next_packet().has_value());  // only one packet was written

    std::fclose(f);
}

TEST_CASE("Reader::link_type reflects the IDB, populated by the time the first packet returns") {
    std::FILE* f = std::tmpfile();
    REQUIRE(f != nullptr);

    Writer writer(f);
    writer.write_section_header();
    writer.write_interface_description(65535, /*link_type=*/12);  // DLT_RAW, arbitrary for this test

    std::vector<unsigned char> data = {0x01};
    writer.write_packet(0, 1, 0, data, 1);

    std::fflush(f);
    std::fseek(f, 0, SEEK_SET);

    Reader reader(f);
    CHECK_FALSE(reader.link_type().has_value());  // nothing read yet
    auto record = reader.next_packet();
    REQUIRE(record.has_value());
    REQUIRE(reader.link_type().has_value());
    CHECK(*reader.link_type() == 12);

    std::fclose(f);
}

TEST_CASE("pcapng writer/reader round-trip multiple packets in order") {
    std::FILE* f = std::tmpfile();
    REQUIRE(f != nullptr);

    Writer writer(f);
    writer.write_section_header();
    writer.write_interface_description(65535, kLinkTypeEthernet);

    for (unsigned char i = 0; i < 5; ++i) {
        std::vector<unsigned char> data = {i};
        writer.write_packet(0, 1700000000 + i, 0, data, 1);
    }
    std::fflush(f);
    std::fseek(f, 0, SEEK_SET);

    Reader reader(f);
    int count = 0;
    while (auto record = reader.next_packet()) {
        REQUIRE(record->data.size() == 1);
        CHECK(record->data[0] == static_cast<unsigned char>(count));
        ++count;
    }
    CHECK(count == 5);

    std::fclose(f);
}

TEST_CASE("pcapng writer pads packet data to a 4-byte boundary without corrupting the next block") {
    std::FILE* f = std::tmpfile();
    REQUIRE(f != nullptr);

    Writer writer(f);
    writer.write_section_header();
    writer.write_interface_description(65535, kLinkTypeEthernet);

    // 3 bytes of packet data forces padding - the case most likely to
    // misalign the following block if the padding math is wrong.
    std::vector<unsigned char> first = {0x01, 0x02, 0x03};
    std::vector<unsigned char> second = {0xAA, 0xBB};
    writer.write_packet(0, 1, 0, first, 3);
    writer.write_packet(0, 2, 0, second, 2);

    std::fflush(f);
    std::fseek(f, 0, SEEK_SET);

    Reader reader(f);
    auto r1 = reader.next_packet();
    REQUIRE(r1.has_value());
    CHECK(r1->data == first);

    auto r2 = reader.next_packet();
    REQUIRE(r2.has_value());
    CHECK(r2->data == second);

    std::fclose(f);
}

TEST_CASE("Reader rejects a block claiming an implausibly large body instead of allocating it") {
    // Found by fuzzing (fuzz/fuzz_pcapng_reader.cpp): total_len is an
    // untrusted 32-bit value straight from the file. A block claiming
    // ~4GB used to be handed straight to `std::vector` before a single
    // body byte was read, OOM-crashing the process on a corrupt or
    // hostile file. This constructs exactly that: a valid-looking
    // block type, followed by a total_len far beyond anything our own
    // writer would ever produce.
    std::FILE* f = std::tmpfile();
    REQUIRE(f != nullptr);

    std::uint8_t block[8];
    block[0] = 0x06; block[1] = 0x00; block[2] = 0x00; block[3] = 0x00;  // EPB
    block[4] = 0xFF; block[5] = 0xFF; block[6] = 0xFF; block[7] = 0x7F;  // total_len ~2GB
    std::fwrite(block, 1, sizeof(block), f);
    std::fflush(f);
    std::fseek(f, 0, SEEK_SET);

    Reader reader(f);
    CHECK_FALSE(reader.next_packet().has_value());  // rejected, not an OOM attempt

    std::fclose(f);
}