#include "screencast/codec/encoder.h"
#include "screencast/network/h264_packetizer.h"
#include "screencast/network/rtp_packet.h"

#include <algorithm>
#include <array>
#include <chrono>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <initializer_list>
#include <optional>
#include <vector>

namespace {

[[noreturn]] void fail(const char* what) {
    std::fprintf(stderr, "test_rtp: FAIL: %s\n", what);
    std::abort();
}

void check(bool condition, const char* what) {
    if (!condition) {
        fail(what);
    }
}

std::vector<std::byte> bytes(std::initializer_list<int> values) {
    std::vector<std::byte> out;
    out.reserve(values.size());
    for (const int value : values) {
        out.push_back(std::byte{static_cast<std::uint8_t>(value)});
    }
    return out;
}

void append_bytes(std::vector<std::byte>& destination, std::initializer_list<int> values) {
    for (const int value : values) {
        destination.push_back(std::byte{static_cast<std::uint8_t>(value)});
    }
}

bool equal_bytes(const std::vector<std::byte>& lhs, const std::vector<std::byte>& rhs) {
    return lhs.size() == rhs.size() && std::equal(lhs.begin(), lhs.end(), rhs.begin());
}

std::uint8_t u8(std::byte value) {
    return std::to_integer<std::uint8_t>(value);
}

// Joins NAL units with 3-byte start codes, the canonical form the
// depacketizer reconstructs.
std::vector<std::byte> annex_b(const std::vector<std::vector<std::byte>>& nals) {
    std::vector<std::byte> out;
    for (const std::vector<std::byte>& nal : nals) {
        append_bytes(out, {0x00, 0x00, 0x01});
        out.insert(out.end(), nal.begin(), nal.end());
    }
    return out;
}

std::vector<std::byte> nal(int header, std::size_t size) {
    std::vector<std::byte> out;
    out.reserve(size);
    out.push_back(std::byte{static_cast<std::uint8_t>(header)});
    for (std::size_t i = 1; i < size; ++i) {
        out.push_back(std::byte{static_cast<std::uint8_t>((i * 7 + 1) & 0xFF)});
    }
    return out;
}

sc::EncodedFrame make_frame(const std::vector<std::byte>& data, std::uint32_t rtp_timestamp) {
    sc::EncodedFrame frame;
    frame.data = data;
    frame.rtp_timestamp = rtp_timestamp;
    frame.is_keyframe = true;
    return frame;
}

sc::RtpPacketizerConfig test_config(std::uint16_t initial_sequence, std::size_t mtu) {
    sc::RtpPacketizerConfig config;
    config.ssrc = 0x12345678;
    config.initial_sequence_number = initial_sequence;
    config.mtu = mtu;
    return config;
}

void test_header_roundtrip() {
    sc::RtpHeader header;
    header.marker = true;
    header.payload_type = 97;
    header.sequence_number = 0xABCD;
    header.timestamp = 0x11223344;
    header.ssrc = 0xDEADBEEF;

    std::array<std::byte, 12> buffer{};
    check(header.serialize(buffer), "header serialize");
    const std::optional<sc::RtpHeader> parsed = sc::RtpHeader::parse(buffer);
    check(parsed.has_value(), "header parse");
    check(parsed->version == 2, "version");
    check(parsed->marker, "marker");
    check(parsed->payload_type == 97, "payload type");
    check(parsed->sequence_number == 0xABCD, "sequence number");
    check(parsed->timestamp == 0x11223344, "timestamp");
    check(parsed->ssrc == 0xDEADBEEF, "ssrc");
    check(!parsed->padding && !parsed->extension && parsed->csrc_count == 0, "flags");
}

void test_header_rejections() {
    std::array<std::byte, 12> buffer{};

    sc::RtpHeader bad_version;
    bad_version.version = 3;
    check(!bad_version.serialize(buffer), "reject version 3");

    sc::RtpHeader bad_csrc;
    bad_csrc.csrc_count = 2;
    check(!bad_csrc.serialize(buffer), "reject csrc count");

    sc::RtpHeader bad_extension;
    bad_extension.extension = true;
    check(!bad_extension.serialize(buffer), "reject extension flag");

    const std::array<std::byte, 12> zeros{}; // version 0 on the wire
    check(!sc::RtpHeader::parse(zeros).has_value(), "reject version 0 input");
}

void test_packet_roundtrip() {
    sc::RtpPacket packet;
    packet.header.sequence_number = 7;
    packet.header.timestamp = 0x0A0B0C0D;
    packet.header.ssrc = 0x01020304;
    packet.payload = bytes({0x67, 0x42, 0x00, 0x01, 0xFF});

    const std::vector<std::byte> wire = packet.serialize();
    check(!wire.empty(), "packet serialize");
    const std::optional<sc::RtpPacket> parsed = sc::RtpPacket::parse(wire);
    check(parsed.has_value(), "packet parse");
    check(parsed->header.sequence_number == 7 && parsed->header.timestamp == 0x0A0B0C0D &&
              parsed->header.ssrc == 0x01020304 && parsed->header.payload_type == 96,
          "packet header fields");
    check(equal_bytes(parsed->payload, packet.payload), "packet payload");
}

void test_packet_parse_tolerances() {
    // CSRC list is skipped, payload preserved.
    std::vector<std::byte> csrc_wire = bytes({0x81, 0x60, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x11, 0x22, 0x33, 0x44});
    append_bytes(csrc_wire, {0xDE, 0xAD, 0xBE, 0xEF});
    append_bytes(csrc_wire, {0xAA, 0xBB});
    const std::optional<sc::RtpPacket> csrc_parsed = sc::RtpPacket::parse(csrc_wire);
    check(csrc_parsed.has_value(), "csrc parse");
    check(csrc_parsed->header.csrc_count == 1, "csrc count");
    check(equal_bytes(csrc_parsed->payload, bytes({0xAA, 0xBB})), "csrc payload preserved");

    // Extension header is skipped, payload preserved.
    std::vector<std::byte> ext_wire = bytes({0x90, 0x60, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x11, 0x22, 0x33, 0x44});
    append_bytes(ext_wire, {0xAB, 0xCD, 0x00, 0x01}); // profile + 1 word
    append_bytes(ext_wire, {0x55, 0x55, 0x55, 0x55}); // extension data
    append_bytes(ext_wire, {0xAA, 0xBB});
    const std::optional<sc::RtpPacket> ext_parsed = sc::RtpPacket::parse(ext_wire);
    check(ext_parsed.has_value(), "extension parse");
    check(ext_parsed->header.extension, "extension flag");
    check(equal_bytes(ext_parsed->payload, bytes({0xAA, 0xBB})), "extension payload preserved");

    // Padding is stripped.
    std::vector<std::byte> pad_wire = bytes({0xA0, 0x60, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x11, 0x22, 0x33, 0x44});
    append_bytes(pad_wire, {0xAA, 0xBB});
    append_bytes(pad_wire, {0x00, 0x00, 0x03}); // 3 padding bytes, count last
    const std::optional<sc::RtpPacket> pad_parsed = sc::RtpPacket::parse(pad_wire);
    check(pad_parsed.has_value(), "padding parse");
    check(pad_parsed->header.padding, "padding flag");
    check(equal_bytes(pad_parsed->payload, bytes({0xAA, 0xBB})), "padding stripped");

    // Malformed input is rejected.
    check(!sc::RtpPacket::parse(bytes({0x80, 0x60, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x11, 0x22, 0x33})).has_value(),
          "reject short packet");

    std::vector<std::byte> truncated_csrc =
        bytes({0x83, 0x60, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x11, 0x22, 0x33, 0x44});
    append_bytes(truncated_csrc, {0xAA, 0xBB});
    check(!sc::RtpPacket::parse(truncated_csrc).has_value(), "reject truncated csrc");

    std::vector<std::byte> overlong_extension =
        bytes({0x90, 0x60, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x11, 0x22, 0x33, 0x44});
    append_bytes(overlong_extension, {0xAB, 0xCD, 0x00, 0x04}); // claims 4 words
    check(!sc::RtpPacket::parse(overlong_extension).has_value(), "reject overlong extension");

    std::vector<std::byte> zero_padding =
        bytes({0xA0, 0x60, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x11, 0x22, 0x33, 0x44});
    append_bytes(zero_padding, {0xAA, 0x00});
    check(!sc::RtpPacket::parse(zero_padding).has_value(), "reject zero padding count");

    std::vector<std::byte> oversized_padding =
        bytes({0xA0, 0x60, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x11, 0x22, 0x33, 0x44});
    append_bytes(oversized_padding, {0x00, 0x09});
    check(!sc::RtpPacket::parse(oversized_padding).has_value(), "reject oversized padding");
}

void test_single_nal_packetization() {
    const std::vector<std::byte> sps = bytes({0x67, 0x42, 0x00, 0x1F});
    const std::vector<std::byte> pps = bytes({0x68, 0xCE, 0x06, 0x0D});
    const sc::EncodedFrame frame = make_frame(annex_b({sps, pps}), 90000);

    sc::H264Packetizer packetizer(test_config(0x0100, 1200));
    const std::vector<sc::RtpPacket> packets = packetizer.packetize(frame);
    check(packets.size() == 2, "two packets");
    check(equal_bytes(packets[0].payload, sps), "sps payload");
    check(equal_bytes(packets[1].payload, pps), "pps payload");
    check(!packets[0].header.marker && packets[1].header.marker, "marker placement");
    check(packets[0].header.sequence_number == 0x0100 && packets[1].header.sequence_number == 0x0101, "sequences");
    check(packets[0].header.timestamp == 90000 && packets[1].header.timestamp == 90000, "timestamps");
    check(packets[0].header.ssrc == 0x12345678, "ssrc");
}

void test_three_byte_start_codes() {
    const std::vector<std::byte> sps = bytes({0x67, 0x01});
    const std::vector<std::byte> pps = bytes({0x68, 0x02});

    std::vector<std::byte> data = bytes({0x00, 0x00, 0x01});
    append_bytes(data, {0x67, 0x01});
    append_bytes(data, {0x00, 0x00, 0x01});
    append_bytes(data, {0x68, 0x02});

    sc::H264Packetizer packetizer(test_config(1, 1200));
    const std::vector<sc::RtpPacket> packets = packetizer.packetize(make_frame(data, 1));
    check(packets.size() == 2, "3-byte start codes split");
    check(equal_bytes(packets[0].payload, sps) && equal_bytes(packets[1].payload, pps), "3-byte start code payloads");
}

void test_four_byte_start_codes() {
    // A zero byte preceding a start code belongs to the previous NAL, so
    // 4-byte start code streams keep their exact byte layout.
    const std::vector<std::byte> sps_with_trailing_zero = bytes({0x67, 0x01, 0x00});
    const std::vector<std::byte> pps = bytes({0x68, 0x02});

    std::vector<std::byte> data = bytes({0x00, 0x00, 0x00, 0x01});
    append_bytes(data, {0x67, 0x01});
    append_bytes(data, {0x00, 0x00, 0x00, 0x01});
    append_bytes(data, {0x68, 0x02});

    sc::H264Packetizer packetizer(test_config(1, 1200));
    const std::vector<sc::RtpPacket> packets = packetizer.packetize(make_frame(data, 1));
    check(packets.size() == 2, "4-byte start codes split");
    check(equal_bytes(packets[0].payload, sps_with_trailing_zero), "absorbed trailing zero");
    check(equal_bytes(packets[1].payload, pps), "4-byte start code last payload");
}

void test_fu_a_fragmentation() {
    const std::vector<std::byte> big_nal = nal(0x65, 26); // IDR slice, 25 payload bytes
    const sc::EncodedFrame frame = make_frame(annex_b({big_nal}), 12345);

    // 24-byte packets: 12 header + 2 FU bytes + 10 data per chunk.
    sc::H264Packetizer packetizer(test_config(0x00F0, 24));
    const std::vector<sc::RtpPacket> packets = packetizer.packetize(frame);
    check(packets.size() == 3, "three fu-a packets");
    for (const sc::RtpPacket& packet : packets) {
        check(packet.serialize().size() <= 24, "mtu respected");
        check(u8(packet.payload[0]) == ((0x65 & 0xE0) | 28), "fu indicator");
        check((u8(packet.payload[1]) & 0x1F) == 5, "nal type preserved");
    }
    check((u8(packets[0].payload[1]) & 0x80) != 0, "start flag");
    check((u8(packets[1].payload[1]) & 0xC0) == 0, "middle flags");
    check((u8(packets[2].payload[1]) & 0x40) != 0, "end flag");
    check(!packets[0].header.marker && !packets[1].header.marker && packets[2].header.marker, "fu marker");
    check(packets[0].header.sequence_number == 0x00F0 && packets[1].header.sequence_number == 0x00F1 &&
              packets[2].header.sequence_number == 0x00F2,
          "fu sequences");

    std::vector<std::byte> reassembled;
    reassembled.push_back(
        std::byte{static_cast<std::uint8_t>((u8(packets[0].payload[0]) & 0xE0) | (u8(packets[0].payload[1]) & 0x1F))});
    for (const sc::RtpPacket& packet : packets) {
        reassembled.insert(reassembled.end(), packet.payload.begin() + 2, packet.payload.end());
    }
    check(equal_bytes(reassembled, big_nal), "fu reassembly");
}

void test_depacketize_roundtrip() {
    const std::vector<std::byte> sps = bytes({0x67, 0x42, 0x00});
    const std::vector<std::byte> pps = bytes({0x68, 0xCE});
    const std::vector<std::byte> big = nal(0x65, 40);
    const std::vector<std::byte> access_unit = annex_b({sps, pps, big});

    // MTU 20: SPS and PPS fit single packets; the 40-byte NAL becomes 7 FU-A
    // chunks of 6 bytes (39 payload bytes), 9 packets total.
    sc::H264Packetizer packetizer(test_config(0x1000, 20));
    const std::vector<sc::RtpPacket> packets = packetizer.packetize(make_frame(access_unit, 3000));
    check(packets.size() == 9, "round-trip packet count");

    sc::H264Depacketizer depacketizer;
    std::optional<std::vector<std::byte>> completed;
    bool any_dropped = false;
    for (const sc::RtpPacket& packet : packets) {
        const sc::DepacketizeResult result = depacketizer.depacketize(packet);
        if (result.access_unit.has_value()) {
            check(!completed.has_value(), "only one completion");
            completed = std::move(result.access_unit);
        }
        any_dropped = any_dropped || result.frame_dropped;
    }
    check(completed.has_value(), "frame completed");
    check(!any_dropped, "no dropped frames in a clean stream");
    check(equal_bytes(*completed, access_unit), "access unit round-trip");
}

void test_depacketizer_drops_gapped_frames() {
    const std::vector<std::byte> big_nal = nal(0x65, 26);
    sc::H264Packetizer packetizer(test_config(0x0100, 24));
    const std::vector<sc::RtpPacket> packets = packetizer.packetize(make_frame(annex_b({big_nal}), 5000));
    check(packets.size() == 3, "gap test packet count");

    sc::H264Depacketizer depacketizer;
    check(!depacketizer.depacketize(packets[0]).access_unit.has_value(), "first fu chunk accepted");
    // packets[1] is lost in transit; the tail cannot complete the frame.
    const sc::DepacketizeResult tail = depacketizer.depacketize(packets[2]);
    check(!tail.access_unit.has_value(), "tail after gap dropped");
    check(tail.frame_dropped, "drop reported after gap");
}

void test_depacketizer_separate_frames() {
    const std::vector<std::byte> f1 = annex_b({bytes({0x67, 0x01})});
    const std::vector<std::byte> f2 = annex_b({bytes({0x41, 0x02})});

    sc::H264Packetizer packetizer(test_config(0x0001, 1200));
    const std::vector<sc::RtpPacket> first = packetizer.packetize(make_frame(f1, 90000));
    const std::vector<sc::RtpPacket> second = packetizer.packetize(make_frame(f2, 90000));

    sc::H264Depacketizer depacketizer;
    const sc::DepacketizeResult au1 = depacketizer.depacketize(first[0]);
    check(au1.access_unit.has_value() && equal_bytes(*au1.access_unit, f1), "first frame");
    check(!au1.frame_dropped, "first frame not dropped");

    // Same RTP timestamp on purpose: the marker alone separates frames.
    const sc::DepacketizeResult au2 = depacketizer.depacketize(second[0]);
    check(au2.access_unit.has_value() && equal_bytes(*au2.access_unit, f2), "second frame with same timestamp");
    check(!au2.frame_dropped, "second frame not dropped");
}

void test_jitter_buffer_in_order() {
    sc::RtpJitterBuffer jitter;
    sc::RtpPacket packet;
    packet.header.sequence_number = 100;

    std::vector<sc::RtpPacket> released = jitter.push(packet);
    check(released.size() == 1 && released[0].header.sequence_number == 100, "in-order releases immediately");

    packet.header.sequence_number = 101;
    released = jitter.push(std::move(packet));
    check(released.size() == 1 && released[0].header.sequence_number == 101, "next packet releases too");
}

void test_jitter_buffer_reorders() {
    sc::RtpJitterBuffer jitter;

    sc::RtpPacket first;
    first.header.sequence_number = 1;
    std::vector<sc::RtpPacket> released = jitter.push(std::move(first));
    check(released.size() == 1 && released[0].header.sequence_number == 1, "first packet releases");

    // Arrives ahead of its predecessor: held, not delivered.
    sc::RtpPacket third;
    third.header.sequence_number = 3;
    released = jitter.push(std::move(third));
    check(released.empty(), "gap holds packets");

    sc::RtpPacket second;
    second.header.sequence_number = 2;
    released = jitter.push(std::move(second));
    check(released.size() == 2, "held packets release in order");
    check(released[0].header.sequence_number == 2 && released[1].header.sequence_number == 3,
          "released in sequence order");
}

void test_jitter_buffer_overflow_and_stragglers() {
    // Small depth: a persistent gap overflows the buffer and flushes what
    // is there, so genuine loss reaches the depacketizer instead of
    // stalling delivery.
    sc::RtpJitterBuffer jitter(4, std::chrono::milliseconds{500});

    sc::RtpPacket packet;
    packet.header.sequence_number = 10;
    check(jitter.push(std::move(packet)).size() == 1, "first releases");

    std::vector<sc::RtpPacket> released;
    for (std::uint16_t sequence = 12; sequence < 17; ++sequence) {
        sc::RtpPacket missing;
        missing.header.sequence_number = sequence;
        for (sc::RtpPacket out : jitter.push(std::move(missing))) {
            released.push_back(std::move(out));
        }
    }
    check(released.size() == 5, "overflow flushes the backlog");
    for (std::size_t i = 0; i < released.size(); ++i) {
        check(released[i].header.sequence_number == 12 + i, "flushed in order");
    }

    // A straggler older than the delivered sequence is discarded, not
    // re-inserted out of order.
    sc::RtpPacket straggler;
    straggler.header.sequence_number = 11;
    check(jitter.push(std::move(straggler)).empty(), "straggler discarded");
}

void test_sequence_wrap() {
    sc::H264Packetizer packetizer(test_config(0xFFFE, 1200));
    const std::vector<sc::RtpPacket> packets =
        packetizer.packetize(make_frame(annex_b({bytes({0x67, 0x01}), bytes({0x68, 0x02}), bytes({0x65, 0x03})}), 100));
    check(packets.size() == 3, "wrap packet count");
    check(packets[0].header.sequence_number == 0xFFFE && packets[1].header.sequence_number == 0xFFFF &&
              packets[2].header.sequence_number == 0x0000,
          "sequence wrap");
}

void test_default_config_randomizes() {
    sc::H264Packetizer first;
    sc::H264Packetizer second;
    const std::vector<sc::RtpPacket> from_first = first.packetize(make_frame(annex_b({bytes({0x67, 0x01})}), 1));
    const std::vector<sc::RtpPacket> from_second = second.packetize(make_frame(annex_b({bytes({0x67, 0x01})}), 1));
    check(!from_first.empty() && !from_second.empty(), "default packetize");
    check(from_first[0].header.ssrc != from_second[0].header.ssrc, "random ssrc");
    check(from_first[0].header.sequence_number != from_second[0].header.sequence_number, "random sequence");
}

void test_empty_inputs() {
    sc::H264Packetizer packetizer(test_config(1, 1200));
    check(packetizer.packetize(sc::EncodedFrame{}).empty(), "empty frame data");
    check(packetizer.packetize(make_frame(bytes({0x00, 0x01, 0x02}), 1)).empty(), "no start codes");

    sc::H264Packetizer tiny_mtu(test_config(1, 8));
    check(tiny_mtu.packetize(make_frame(annex_b({bytes({0x67, 0x01})}), 1)).empty(), "unusable mtu");
}

} // namespace

int main() {
    test_header_roundtrip();
    test_header_rejections();
    test_packet_roundtrip();
    test_packet_parse_tolerances();
    test_single_nal_packetization();
    test_three_byte_start_codes();
    test_four_byte_start_codes();
    test_fu_a_fragmentation();
    test_depacketize_roundtrip();
    test_depacketizer_drops_gapped_frames();
    test_depacketizer_separate_frames();
    test_jitter_buffer_in_order();
    test_jitter_buffer_reorders();
    test_jitter_buffer_overflow_and_stragglers();
    test_sequence_wrap();
    test_default_config_randomizes();
    test_empty_inputs();
    std::puts("test_rtp: all checks passed");
    return 0;
}