#include "media_reassembler.h"

#include <chrono>
#include <cstdint>
#include <iostream>
#include <string>

namespace {

#define CHECK_TRUE(expression)                                                \
    do {                                                                      \
        if (!(expression)) {                                                  \
            std::cerr << "CHECK failed at line " << __LINE__ << ": "         \
                      << #expression << '\n';                                 \
            return false;                                                     \
        }                                                                     \
    } while (false)

using cmvr::test::quic_gateway::MediaReassembler;
using cmvr::test::quic_gateway::ReassemblyResult;
using Clock = std::chrono::steady_clock;

cmvr::quic_edge::DatagramHeader header(
    const std::uint16_t fragment_index,
    const std::uint32_t offset,
    const std::uint16_t payload_size,
    const std::uint64_t frame_sequence = 1U)
{
    cmvr::quic_edge::DatagramHeader value;
    value.kind = cmvr::media::MediaKind::VIDEO;
    value.fragment_index = fragment_index;
    value.fragment_count = 2U;
    value.payload_size = payload_size;
    value.track_id = 7U;
    value.codec_generation = 3U;
    value.session_epoch = 11U;
    value.packet_sequence = fragment_index;
    value.frame_sequence = frame_sequence;
    value.capture_timestamp_us = 100U;
    value.frame_size = 6U;
    value.fragment_offset = offset;
    return value;
}

bool completesOutOfOrderAndDetectsDuplicates()
{
    MediaReassembler reassembler(
        1024U, 8U, 1024U, std::chrono::milliseconds(100));
    CHECK_TRUE(reassembler.reset(11U) == 0U);
    const auto now = Clock::now();
    const std::string second = "def";
    auto result = reassembler.accept(
        header(1U, 3U, 3U),
        reinterpret_cast<const std::uint8_t*>(second.data()),
        second.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::ACCEPTED);

    result = reassembler.accept(
        header(1U, 3U, 3U),
        reinterpret_cast<const std::uint8_t*>(second.data()),
        second.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::DUPLICATE);

    const std::string first = "abc";
    result = reassembler.accept(
        header(0U, 0U, 3U),
        reinterpret_cast<const std::uint8_t*>(first.data()),
        first.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::COMPLETED);
    CHECK_TRUE(result.completed_frame.has_value());
    CHECK_TRUE(std::string(
        result.completed_frame->payload.begin(),
        result.completed_frame->payload.end()) == "abcdef");
    CHECK_TRUE(reassembler.inFlightFrames() == 0U);
    CHECK_TRUE(reassembler.bufferedBytes() == 0U);
    return true;
}

bool rejectsOverlappingFragments()
{
    MediaReassembler reassembler(
        1024U, 8U, 1024U, std::chrono::milliseconds(100));
    reassembler.reset(11U);
    const auto now = Clock::now();
    const std::string first = "abcd";
    auto first_header = header(0U, 0U, 4U);
    auto result = reassembler.accept(
        first_header,
        reinterpret_cast<const std::uint8_t*>(first.data()),
        first.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::ACCEPTED);

    const std::string second = "def";
    result = reassembler.accept(
        header(1U, 3U, 3U),
        reinterpret_cast<const std::uint8_t*>(second.data()),
        second.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::INVALID);
    CHECK_TRUE(reassembler.inFlightFrames() == 0U);
    return true;
}

bool rejectsConflictingFlags()
{
    MediaReassembler reassembler(
        1024U, 8U, 1024U, std::chrono::milliseconds(100));
    reassembler.reset(11U);
    const auto now = Clock::now();
    const std::string payload = "abc";

    auto first_header = header(0U, 0U, 3U);
    first_header.flags = 1U;
    auto result = reassembler.accept(
        first_header,
        reinterpret_cast<const std::uint8_t*>(payload.data()),
        payload.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::ACCEPTED);

    auto second_header = header(1U, 3U, 3U);
    second_header.flags = 0U;
    result = reassembler.accept(
        second_header,
        reinterpret_cast<const std::uint8_t*>(payload.data()),
        payload.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::INVALID);
    CHECK_TRUE(reassembler.inFlightFrames() == 0U);
    CHECK_TRUE(reassembler.bufferedBytes() == 0U);
    return true;
}

bool rejectsConflictingCaptureTimestamp()
{
    MediaReassembler reassembler(
        1024U, 8U, 1024U, std::chrono::milliseconds(100));
    reassembler.reset(11U);
    const auto now = Clock::now();
    const std::string payload = "abc";

    auto result = reassembler.accept(
        header(0U, 0U, 3U),
        reinterpret_cast<const std::uint8_t*>(payload.data()),
        payload.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::ACCEPTED);

    auto second_header = header(1U, 3U, 3U);
    second_header.capture_timestamp_us = 101U;
    result = reassembler.accept(
        second_header,
        reinterpret_cast<const std::uint8_t*>(payload.data()),
        payload.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::INVALID);
    CHECK_TRUE(reassembler.inFlightFrames() == 0U);
    CHECK_TRUE(reassembler.bufferedBytes() == 0U);
    return true;
}

bool expiresAndBoundsIncompleteFrames()
{
    MediaReassembler expiring(
        1024U, 8U, 1024U, std::chrono::milliseconds(10));
    expiring.reset(11U);
    const auto now = Clock::now();
    const std::string payload = "abc";
    auto result = expiring.accept(
        header(0U, 0U, 3U, 1U),
        reinterpret_cast<const std::uint8_t*>(payload.data()),
        payload.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::ACCEPTED);
    result = expiring.accept(
        header(0U, 0U, 3U, 2U),
        reinterpret_cast<const std::uint8_t*>(payload.data()),
        payload.size(), now + std::chrono::milliseconds(11));
    CHECK_TRUE(result.expired_frames == 1U);

    MediaReassembler bounded(
        4U, 8U, 1024U, std::chrono::milliseconds(100));
    bounded.reset(11U);
    result = bounded.accept(
        header(0U, 0U, 3U),
        reinterpret_cast<const std::uint8_t*>(payload.data()),
        payload.size(), now);
    CHECK_TRUE(result.status == ReassemblyResult::Status::ACCEPTED);
    result = bounded.accept(
        header(1U, 3U, 3U),
        reinterpret_cast<const std::uint8_t*>(payload.data()),
        payload.size(), now);
    CHECK_TRUE(
        result.status == ReassemblyResult::Status::CAPACITY_DROPPED);
    CHECK_TRUE(bounded.inFlightFrames() == 0U);
    return true;
}

} // namespace

int main()
{
    if (!completesOutOfOrderAndDetectsDuplicates()) return 1;
    if (!rejectsOverlappingFragments()) return 1;
    if (!rejectsConflictingFlags()) return 1;
    if (!rejectsConflictingCaptureTimestamp()) return 1;
    if (!expiresAndBoundsIncompleteFrames()) return 1;
    std::cout << "cmvr_quic_media_reassembler_test passed\n";
    return 0;
}