#ifndef CMVR_ES_TEST_QUIC_GATEWAY_MEDIA_REASSEMBLER_H
#define CMVR_ES_TEST_QUIC_GATEWAY_MEDIA_REASSEMBLER_H

#include <chrono>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <unordered_map>
#include <vector>

#include "service/quic_edge/include/quic_edge_types.h"

namespace cmvr::test::quic_gateway {

struct ReassembledFrame {
    cmvr::quic_edge::DatagramHeader header;
    std::vector<std::uint8_t> payload;
};

struct ReassemblyResult {
    enum class Status {
        ACCEPTED,
        DUPLICATE,
        COMPLETED,
        INVALID,
        CAPACITY_DROPPED,
    };

    Status status{Status::INVALID};
    std::optional<ReassembledFrame> completed_frame;
    std::uint64_t expired_frames{0};
    std::uint64_t evicted_frames{0};
};

class MediaReassembler {
public:
    MediaReassembler(std::size_t maximum_bytes,
                     std::size_t maximum_frames,
                     std::size_t maximum_frame_bytes,
                     std::chrono::milliseconds timeout);

    std::uint64_t reset(std::uint64_t session_epoch);

    ReassemblyResult accept(
        const cmvr::quic_edge::DatagramHeader& header,
        const std::uint8_t* payload,
        std::size_t payload_size,
        std::chrono::steady_clock::time_point now);

    std::uint64_t clear();
    std::size_t bufferedBytes() const { return buffered_bytes_; }
    std::size_t inFlightFrames() const { return frames_.size(); }
    std::uint64_t sessionEpoch() const { return session_epoch_; }

private:
    struct FrameKey {
        std::uint64_t session_epoch{0};
        std::uint32_t track_id{0};
        std::uint64_t frame_sequence{0};

        bool operator==(const FrameKey& other) const
        {
            return session_epoch == other.session_epoch &&
                   track_id == other.track_id &&
                   frame_sequence == other.frame_sequence;
        }
    };

    struct FrameKeyHash {
        std::size_t operator()(const FrameKey& key) const;
    };

    struct Fragment {
        std::uint32_t offset{0};
        std::vector<std::uint8_t> payload;
    };

    struct PartialFrame {
        cmvr::quic_edge::DatagramHeader first_header;
        std::unordered_map<std::uint16_t, Fragment> fragments;
        std::size_t received_bytes{0};
        std::chrono::steady_clock::time_point created_at;
        std::chrono::steady_clock::time_point updated_at;
    };

    std::uint64_t expire(std::chrono::steady_clock::time_point now);
    std::uint64_t evictOldest();
    std::uint64_t evictOldestExcept(const FrameKey& protected_key);
    void eraseFrame(
        std::unordered_map<FrameKey, PartialFrame, FrameKeyHash>::iterator it);

    std::size_t maximum_bytes_;
    std::size_t maximum_frames_;
    std::size_t maximum_frame_bytes_;
    std::chrono::milliseconds timeout_;
    std::uint64_t session_epoch_{0};
    std::size_t buffered_bytes_{0};
    std::unordered_map<FrameKey, PartialFrame, FrameKeyHash> frames_;
};

} // namespace cmvr::test::quic_gateway

#endif // CMVR_ES_TEST_QUIC_GATEWAY_MEDIA_REASSEMBLER_H