#include "media_reassembler.h"

#include <algorithm>
#include <utility>

namespace cmvr::test::quic_gateway {
namespace {

constexpr std::size_t kMaximumFragmentsPerFrame = 8192U;

bool rangesOverlap(const std::uint32_t left_offset,
                   const std::size_t left_size,
                   const std::uint32_t right_offset,
                   const std::size_t right_size)
{
    const std::uint64_t left_end =
        static_cast<std::uint64_t>(left_offset) + left_size;
    const std::uint64_t right_end =
        static_cast<std::uint64_t>(right_offset) + right_size;
    return static_cast<std::uint64_t>(left_offset) < right_end &&
           static_cast<std::uint64_t>(right_offset) < left_end;
}

} // namespace

MediaReassembler::MediaReassembler(
    const std::size_t maximum_bytes,
    const std::size_t maximum_frames,
    const std::size_t maximum_frame_bytes,
    const std::chrono::milliseconds timeout)
    : maximum_bytes_(std::max<std::size_t>(1U, maximum_bytes)),
      maximum_frames_(std::max<std::size_t>(1U, maximum_frames)),
      maximum_frame_bytes_(std::max<std::size_t>(1U, maximum_frame_bytes)),
      timeout_(std::max(std::chrono::milliseconds(1), timeout))
{
}

std::size_t MediaReassembler::FrameKeyHash::operator()(
    const FrameKey& key) const
{
    std::size_t value = std::hash<std::uint64_t>{}(key.session_epoch);
    value ^= std::hash<std::uint32_t>{}(key.track_id) +
             0x9e3779b9U + (value << 6U) + (value >> 2U);
    value ^= std::hash<std::uint64_t>{}(key.frame_sequence) +
             0x9e3779b9U + (value << 6U) + (value >> 2U);
    return value;
}

std::uint64_t MediaReassembler::reset(const std::uint64_t session_epoch)
{
    const std::uint64_t dropped = clear();
    session_epoch_ = session_epoch;
    return dropped;
}

ReassemblyResult MediaReassembler::accept(
    const cmvr::quic_edge::DatagramHeader& header,
    const std::uint8_t* payload,
    const std::size_t payload_size,
    const std::chrono::steady_clock::time_point now)
{
    ReassemblyResult result;
    result.expired_frames = expire(now);

    if (!payload || payload_size == 0U ||
        header.session_epoch == 0U ||
        header.session_epoch != session_epoch_ ||
        header.payload_size != payload_size ||
        header.frame_size == 0U ||
        header.frame_size > maximum_frame_bytes_ ||
        header.fragment_count == 0U ||
        header.fragment_count > kMaximumFragmentsPerFrame ||
        header.fragment_index >= header.fragment_count ||
        header.fragment_offset > header.frame_size ||
        payload_size > header.frame_size - header.fragment_offset) {
        result.status = ReassemblyResult::Status::INVALID;
        return result;
    }

    const FrameKey key{
        header.session_epoch, header.track_id, header.frame_sequence};
    auto frame_it = frames_.find(key);
    if (frame_it == frames_.end()) {
        while (frames_.size() >= maximum_frames_) {
            result.evicted_frames += evictOldest();
        }
        PartialFrame frame;
        frame.first_header = header;
        frame.created_at = now;
        frame.updated_at = now;
        try {
            frame_it = frames_.emplace(key, std::move(frame)).first;
        } catch (...) {
            result.status = ReassemblyResult::Status::CAPACITY_DROPPED;
            return result;
        }
    }

    PartialFrame& frame = frame_it->second;
    if (frame.first_header.fragment_count != header.fragment_count ||
        frame.first_header.frame_size != header.frame_size ||
        frame.first_header.kind != header.kind ||
        frame.first_header.codec_generation != header.codec_generation ||
        frame.first_header.flags != header.flags ||
        frame.first_header.capture_timestamp_us !=
            header.capture_timestamp_us) {
        eraseFrame(frame_it);
        result.status = ReassemblyResult::Status::INVALID;
        return result;
    }

    const auto duplicate = frame.fragments.find(header.fragment_index);
    if (duplicate != frame.fragments.end()) {
        const Fragment& existing = duplicate->second;
        const bool identical =
            existing.offset == header.fragment_offset &&
            existing.payload.size() == payload_size &&
            std::equal(existing.payload.begin(), existing.payload.end(), payload);
        result.status = identical
            ? ReassemblyResult::Status::DUPLICATE
            : ReassemblyResult::Status::INVALID;
        if (!identical) eraseFrame(frame_it);
        return result;
    }

    for (const auto& [index, existing] : frame.fragments) {
        (void)index;
        if (rangesOverlap(header.fragment_offset, payload_size,
                          existing.offset, existing.payload.size())) {
            eraseFrame(frame_it);
            result.status = ReassemblyResult::Status::INVALID;
            return result;
        }
    }

    if (payload_size > maximum_bytes_) {
        eraseFrame(frame_it);
        result.status = ReassemblyResult::Status::CAPACITY_DROPPED;
        return result;
    }
    while (buffered_bytes_ > maximum_bytes_ - payload_size) {
        const std::uint64_t evicted = evictOldestExcept(key);
        if (evicted == 0U) {
            frame_it = frames_.find(key);
            if (frame_it != frames_.end()) eraseFrame(frame_it);
            result.status = ReassemblyResult::Status::CAPACITY_DROPPED;
            return result;
        }
        result.evicted_frames += evicted;
    }

    Fragment fragment;
    fragment.offset = header.fragment_offset;
    try {
        fragment.payload.assign(payload, payload + payload_size);
        frame.fragments.emplace(header.fragment_index, std::move(fragment));
    } catch (...) {
        eraseFrame(frame_it);
        result.status = ReassemblyResult::Status::CAPACITY_DROPPED;
        return result;
    }
    frame.received_bytes += payload_size;
    frame.updated_at = now;
    buffered_bytes_ += payload_size;

    if (frame.fragments.size() != header.fragment_count ||
        frame.received_bytes != header.frame_size) {
        result.status = ReassemblyResult::Status::ACCEPTED;
        return result;
    }

    ReassembledFrame complete;
    complete.header = frame.first_header;
    try {
        complete.payload.resize(header.frame_size);
    } catch (...) {
        eraseFrame(frame_it);
        result.status = ReassemblyResult::Status::CAPACITY_DROPPED;
        return result;
    }
    for (const auto& [index, value] : frame.fragments) {
        (void)index;
        if (value.offset > complete.payload.size() ||
            value.payload.size() > complete.payload.size() - value.offset) {
            eraseFrame(frame_it);
            result.status = ReassemblyResult::Status::INVALID;
            return result;
        }
        std::copy(value.payload.begin(), value.payload.end(),
                  complete.payload.begin() + value.offset);
    }

    eraseFrame(frame_it);
    result.status = ReassemblyResult::Status::COMPLETED;
    result.completed_frame = std::move(complete);
    return result;
}

std::uint64_t MediaReassembler::clear()
{
    const std::uint64_t dropped = frames_.size();
    frames_.clear();
    buffered_bytes_ = 0U;
    return dropped;
}

std::uint64_t MediaReassembler::expire(
    const std::chrono::steady_clock::time_point now)
{
    std::uint64_t expired = 0U;
    for (auto it = frames_.begin(); it != frames_.end();) {
        if (now - it->second.updated_at >= timeout_) {
            buffered_bytes_ -= it->second.received_bytes;
            it = frames_.erase(it);
            ++expired;
        } else {
            ++it;
        }
    }
    return expired;
}

std::uint64_t MediaReassembler::evictOldest()
{
    if (frames_.empty()) return 0U;
    auto oldest = frames_.begin();
    for (auto it = std::next(frames_.begin()); it != frames_.end(); ++it) {
        if (it->second.updated_at < oldest->second.updated_at) oldest = it;
    }
    eraseFrame(oldest);
    return 1U;
}

std::uint64_t MediaReassembler::evictOldestExcept(
    const FrameKey& protected_key)
{
    auto oldest = frames_.end();
    for (auto it = frames_.begin(); it != frames_.end(); ++it) {
        if (it->first == protected_key) continue;
        if (oldest == frames_.end() ||
            it->second.updated_at < oldest->second.updated_at) {
            oldest = it;
        }
    }
    if (oldest == frames_.end()) return 0U;
    eraseFrame(oldest);
    return 1U;
}

void MediaReassembler::eraseFrame(
    std::unordered_map<FrameKey, PartialFrame, FrameKeyHash>::iterator it)
{
    if (it == frames_.end()) return;
    buffered_bytes_ -= it->second.received_bytes;
    frames_.erase(it);
}

} // namespace cmvr::test::quic_gateway