#include <catch2/catch_all.hpp>
#include <cstdint>
#include <cstring>
#include <arpa/inet.h> // htons, htonl
#include "../fpga/hls_simulation/hls_cores.h"
#include "../fpga/hls/ptp.cpp"

// Minimal 64-byte frame layout for a single-beat 512-bit AXI packet.
// We only care about:
//  - Ethernet header (14 bytes) as padding
//  - PTP common header fields: messageType nibble, flagField
//  - preciseOriginTimestamp: ns(32) + seconds(48)
// Other fields are left as zero.
#pragma pack(push, 1)
struct EthernetPTP512 {
    uint8_t eth_pad[14];    // Ethernet dest/src/type (unused in tests)

    // PTP common header (we only fill what we need)
    uint8_t  ts_msgType;    // transportSpecific (upper nibble) | messageType (lower nibble)
    uint8_t  versionPTP;    // not used, keep zero
    uint16_t messageLength; // not used
    uint8_t  domainNumber;  // not used
    uint8_t  reserved1;     // not used
    uint16_t flagField;     // network order
    uint64_t correctionField; // not used
    uint32_t reserved2;     // not used
    uint8_t  sourcePortIdentity[10]; // not used
    uint16_t sequenceId;    // not used
    uint8_t  controlField;  // not used
    int8_t   logMessageInterval; // not used

    // payload for Sync/Follow_Up we care about: preciseOriginTimestamp
    // Offset chosen to match your parser: ns first (32-bit), then seconds (48-bit).
    uint32_t preciseOriginTimestamp_ns_be; // network order
    uint8_t  preciseOriginTimestamp_sec_be[6]; // 48-bit, big-endian
};
#pragma pack(pop)

// Helper: host-to-big-endian for 48-bit seconds into a 6-byte array.
static inline void htobe48(uint64_t host_val, uint8_t out[6]) {
    // Keep low 48 bits
    uint64_t v = host_val & 0x0000FFFFFFFFFFFFULL;
    out[0] = static_cast<uint8_t>((v >> 40) & 0xFF);
    out[1] = static_cast<uint8_t>((v >> 32) & 0xFF);
    out[2] = static_cast<uint8_t>((v >> 24) & 0xFF);
    out[3] = static_cast<uint8_t>((v >> 16) & 0xFF);
    out[4] = static_cast<uint8_t>((v >> 8)  & 0xFF);
    out[5] = static_cast<uint8_t>(v & 0xFF);
}

// Build a 512-bit single-beat AXI PTP packet via a C struct, then cast to ap_uint<512>.
static packet_512_t make_ptp_packet(uint8_t message_type, bool two_step, uint64_t sec48, uint32_t ns32) {
    ap_uint<512> word = 0;

    auto frame = reinterpret_cast<EthernetPTP512 *>(&word);

    // Byte 0 of PTP: transportSpecific=0, messageType in low nibble
    frame->ts_msgType = (0x0 << 4) | (message_type & 0x0F);

    // twoStepFlag: bit 1 (0x0002) in flagField (big-endian)
    uint16_t flags = two_step ? 0x0002 : 0x0000;
    frame->flagField = htons(flags);

    // preciseOriginTimestamp: ns (32-bit BE), seconds (48-bit BE)
    frame->preciseOriginTimestamp_ns_be = htonl(ns32);
    htobe48(sec48, frame->preciseOriginTimestamp_sec_be);

    packet_512_t pkt{};
    pkt.data = word;
    pkt.last = 1;
    return pkt;
}

static void run_cycles(AXI_STREAM& s, int cycles,
                       volatile ap_uint<80>& counter_ns,
                       volatile ap_uint<32>& counter_error,
                       volatile ap_uint<1>& synced_to_ptp) {
    for (int i = 0; i < cycles; ++i) {
        ptp(s, counter_ns, counter_error, synced_to_ptp);
    }
}

TEST_CASE("HLS_PTP_OneStep") {
    AXI_STREAM eth_in;
    volatile ap_uint<80> counter_ns = 0;
    volatile ap_uint<32> counter_error = 0;
    volatile ap_uint<1>  synced_to_ptp = 0;

    run_cycles(eth_in, 10, counter_ns, counter_error, synced_to_ptp);

    uint64_t sec = 56789;
    uint32_t ns  = 999999999U;
    auto pkt = make_ptp_packet(/*message_type=*/0x0, /*two_step=*/false, sec, ns);
    eth_in.write(pkt);

    run_cycles(eth_in, 1, counter_ns, counter_error, synced_to_ptp);

    ap_uint<80> combined = counter_ns; // upper 48 bits = sec, lower 32 bits = ns
    uint64_t sec_read = (uint64_t)(combined >> 32);
    uint32_t ns_read  = (uint32_t)(combined & 0xFFFFFFFFu);

    REQUIRE(synced_to_ptp.read() == 1);
    CHECK(sec_read == sec);
    CHECK(ns_read == ns);
}

TEST_CASE("HLS_PTP_OneStep_incr_1sec") {
    AXI_STREAM eth_in;
    volatile ap_uint<80> counter_ns = 0;
    volatile ap_uint<32> counter_error = 0;
    volatile ap_uint<1>  synced_to_ptp = 0;

    run_cycles(eth_in, 10, counter_ns, counter_error, synced_to_ptp);

    uint64_t sec = 56789;
    uint32_t ns  = 999999999U;
    auto pkt = make_ptp_packet(/*message_type=*/0x0, /*two_step=*/false, sec, ns);
    eth_in.write(pkt);

    run_cycles(eth_in, 2, counter_ns, counter_error, synced_to_ptp);

    ap_uint<80> combined = counter_ns; // upper 48 bits = sec, lower 32 bits = ns
    uint64_t sec_read = (uint64_t)(combined >> 32);
    uint32_t ns_read  = (uint32_t)(combined & 0xFFFFFFFFu);

    REQUIRE(synced_to_ptp.read() == 1);
    CHECK(sec_read == 56789 + 1);
    CHECK(ns_read == 4);
}

TEST_CASE("HLS_PTP_TwoStep") {
    AXI_STREAM eth_in;
    volatile ap_uint<80> counter_ns = 0;
    volatile ap_uint<32> counter_error = 0;
    volatile ap_uint<1>  synced_to_ptp = 0;

    run_cycles(eth_in, 5, counter_ns, counter_error, synced_to_ptp);
    ap_uint<80> before = counter_ns;

    // Two-step Sync
    auto sync_pkt = make_ptp_packet(/*message_type=*/0x0, /*two_step=*/true, /*sec*/5, /*ns*/500);
    eth_in.write(sync_pkt);
    run_cycles(eth_in, 1, counter_ns, counter_error, synced_to_ptp);

    ap_uint<80> after_sync = counter_ns;
    REQUIRE(after_sync > before); // clock advanced by tick, but no snap to sec=5
    uint64_t sec_sync = (uint64_t)(after_sync >> 32);
    REQUIRE(sec_sync != 5);

    // Follow_Up updates to its timestamp
    uint64_t fu_sec = 123456;
    uint32_t fu_ns  = 3456;
    auto follow_up_pkt = make_ptp_packet(/*message_type=*/0x8, /*two_step=*/true, fu_sec, fu_ns);
    eth_in.write(follow_up_pkt);
    run_cycles(eth_in, 1, counter_ns, counter_error, synced_to_ptp);

    ap_uint<80> combined = counter_ns;
    uint64_t sec_read = (uint64_t)(combined >> 32);
    uint32_t ns_read  = (uint32_t)(combined & 0xFFFFFFFFu);

    REQUIRE(synced_to_ptp.read() == 1);
    CHECK(sec_read == fu_sec);
    CHECK(ns_read == fu_ns);
}

TEST_CASE("HLS_PTP_OneStep_error") {
    AXI_STREAM eth_in;
    volatile ap_uint<80> counter_ns = 0;
    volatile ap_uint<32> counter_error = 0;
    volatile ap_uint<1>  synced_to_ptp = 0;

    run_cycles(eth_in, 10, counter_ns, counter_error, synced_to_ptp);

    uint64_t sec = 100;
    uint32_t ns  = 100;
    auto pkt = make_ptp_packet(/*message_type=*/0x0, /*two_step=*/false, sec, ns);
    eth_in.write(pkt);

    run_cycles(eth_in, 11, counter_ns, counter_error, synced_to_ptp);

    ap_uint<80> combined = counter_ns; // upper 48 bits = sec, lower 32 bits = ns
    uint64_t sec_read = (uint64_t)(combined >> 32);
    uint32_t ns_read  = (uint32_t)(combined & 0xFFFFFFFFu);

    REQUIRE(synced_to_ptp.read() == 1);
    CHECK(sec_read == 100);
    CHECK(ns_read == 100 + 5 * 10);

    ns = 100 + 5 * 11 - 7;
    pkt = make_ptp_packet(/*message_type=*/0x0, /*two_step=*/false, sec, ns);
    eth_in.write(pkt);

    run_cycles(eth_in, 1, counter_ns, counter_error, synced_to_ptp);
    CHECK(counter_error.read() == 7);
    combined = counter_ns; // upper 48 bits = sec, lower 32 bits = ns
    sec_read = (uint64_t)(combined >> 32);
    ns_read  = (uint32_t)(combined & 0xFFFFFFFFu);

    CHECK(sec_read == 100);
    CHECK(ns_read == ns);
}

TEST_CASE("HLS_PTP_2s_out_of_sync") {
    AXI_STREAM eth_in;
    volatile ap_uint<80> counter_ns = 0;
    volatile ap_uint<32> counter_error = 0;
    volatile ap_uint<1>  synced_to_ptp = 0;

    auto pkt = make_ptp_packet(/*message_type=*/0x0, /*two_step=*/false, /*sec*/100, /*ns*/0);
    eth_in.write(pkt);
    run_cycles(eth_in, 1, counter_ns, counter_error, synced_to_ptp);
    REQUIRE(synced_to_ptp.read() == 1);

    const uint64_t ticks_2s = (2ULL * 1000000000ULL) / 5ULL;
    run_cycles(eth_in, (int)(ticks_2s + 2), counter_ns, counter_error, synced_to_ptp);

    REQUIRE(synced_to_ptp.read() == 0);
}