package com.cmvr.quic.gateway.protocol;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import org.junit.Assert;
import org.junit.Test;

import java.nio.charset.StandardCharsets;

public class MediaReassemblerTest {

    @Test
    public void shouldDecodeHeaderAndReassembleOutOfOrderFragments() {
        byte[] payload = "cmvr-quic-media-frame".getBytes(StandardCharsets.UTF_8);
        MediaReassembler reassembler = new MediaReassembler(4096, 8, 1024, 2000);
        reassembler.reset(88L);

        ByteBuf second = datagram(2, 0, 1, 2, 5, 7, 88, 2, 9, 1000,
                payload.length, 8, payload, 8, payload.length - 8);
        ByteBuf first = datagram(2, 0, 0, 2, 5, 7, 88, 1, 9, 1000,
                payload.length, 0, payload, 0, 8);
        try {
            DatagramHeader secondHeader = DatagramHeader.decode(second);
            Assert.assertNull(reassembler.accept(secondHeader, second, 10));
            DatagramHeader firstHeader = DatagramHeader.decode(first);
            MediaReassembler.CompletedFrame completed = reassembler.accept(firstHeader, first, 11);
            Assert.assertNotNull(completed);
            Assert.assertArrayEquals(payload, completed.getPayload());
            Assert.assertEquals(0, reassembler.inFlightFrames());
            Assert.assertEquals(0, reassembler.allocatedBytes());
        } finally {
            second.release();
            first.release();
        }
    }

    @Test
    public void shouldExpireIncompleteFrameAndReleaseCapacity() {
        byte[] payload = new byte[] {1, 2, 3, 4};
        MediaReassembler reassembler = new MediaReassembler(64, 2, 32, 10);
        reassembler.reset(1);
        ByteBuf datagram = datagram(1, 1, 0, 2, 1, 1, 1, 1, 1, 1,
                8, 0, payload, 0, payload.length);
        try {
            reassembler.accept(DatagramHeader.decode(datagram), datagram, 100);
            Assert.assertEquals(1, reassembler.inFlightFrames());
            Assert.assertEquals(1, reassembler.expire(110));
            Assert.assertEquals(0, reassembler.allocatedBytes());
        } finally {
            datagram.release();
        }
    }

    private static ByteBuf datagram(int kind, int flags, int fragmentIndex,
                                    int fragmentCount, int trackId, int generation,
                                    long epoch, long packetSequence, long frameSequence,
                                    long captureTimestampUs, int frameSize, int fragmentOffset,
                                    byte[] payload, int payloadOffset, int payloadLength) {
        ByteBuf result = Unpooled.buffer(DatagramHeader.HEADER_BYTES + payloadLength);
        result.writeInt(DatagramHeader.MAGIC);
        result.writeByte(DatagramHeader.PROTOCOL_VERSION);
        result.writeByte(kind);
        result.writeShort(flags);
        result.writeShort(DatagramHeader.HEADER_BYTES);
        result.writeShort(fragmentIndex);
        result.writeShort(fragmentCount);
        result.writeShort(payloadLength);
        result.writeInt(trackId);
        result.writeInt(generation);
        result.writeLong(epoch);
        result.writeLong(packetSequence);
        result.writeLong(frameSequence);
        result.writeLong(captureTimestampUs);
        result.writeInt(frameSize);
        result.writeInt(fragmentOffset);
        result.writeBytes(payload, payloadOffset, payloadLength);
        return result;
    }
}
