package com.cmvr.quic.gateway.protocol;

import io.netty.buffer.ByteBuf;

import java.util.BitSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * 按(sessionEpoch, trackId, frameSequence)重组媒体分片。
 *
 * <p>内存、在途帧数、单帧大小和超时时间全部有硬上限。容量不足时淘汰最早更新的
 * 不完整帧，避免异常终端或丢包造成堆内存持续增长。</p>
 */
public final class MediaReassembler {

    private final long maximumBytes;
    private final int maximumFrames;
    private final int maximumFrameBytes;
    private final long timeoutMs;
    private final LinkedHashMap<FrameKey, PartialFrame> frames =
            new LinkedHashMap<FrameKey, PartialFrame>(16, 0.75F, true);

    private long sessionEpoch;
    private long allocatedBytes;

    public MediaReassembler(long maximumBytes, int maximumFrames,
                            int maximumFrameBytes, long timeoutMs) {
        if (maximumBytes < maximumFrameBytes || maximumFrames <= 0
                || maximumFrameBytes <= 0 || timeoutMs <= 0) {
            throw new IllegalArgumentException("媒体重组容量配置不合法");
        }
        this.maximumBytes = maximumBytes;
        this.maximumFrames = maximumFrames;
        this.maximumFrameBytes = maximumFrameBytes;
        this.timeoutMs = timeoutMs;
    }

    public void reset(long newSessionEpoch) {
        clear();
        sessionEpoch = newSessionEpoch;
    }

    public CompletedFrame accept(DatagramHeader header, ByteBuf datagram, long nowMs) {
        expire(nowMs);
        validate(header);

        FrameKey key = new FrameKey(header.getTrackId(), header.getFrameSequence());
        PartialFrame frame = frames.get(key);
        if (frame == null) {
            reserve(header.getFrameSize());
            frame = new PartialFrame(header, nowMs);
            frames.put(key, frame);
            allocatedBytes += header.getFrameSize();
        } else if (!frame.matches(header)) {
            remove(key, frame);
            throw new IllegalArgumentException("同一媒体帧的分片元数据不一致");
        }

        int index = header.getFragmentIndex();
        int offset = header.getFragmentOffset();
        int length = header.getPayloadSize();
        int payloadIndex = datagram.readerIndex() + DatagramHeader.HEADER_BYTES;
        if (frame.received.get(index)) {
            if (frame.offsets[index] != offset || frame.lengths[index] != length
                    || !equalsBytes(datagram, payloadIndex, frame.payload, offset, length)) {
                remove(key, frame);
                throw new IllegalArgumentException("媒体帧收到冲突的重复分片");
            }
            return null;
        }
        if (overlaps(frame, offset, length)) {
            remove(key, frame);
            throw new IllegalArgumentException("媒体帧分片字节范围重叠");
        }

        datagram.getBytes(payloadIndex, frame.payload, offset, length);
        frame.received.set(index);
        frame.offsets[index] = offset;
        frame.lengths[index] = length;
        frame.receivedBytes += length;
        frame.updatedAtMs = nowMs;

        if (frame.received.cardinality() != header.getFragmentCount()) {
            return null;
        }
        if (frame.receivedBytes != header.getFrameSize()) {
            remove(key, frame);
            throw new IllegalArgumentException("媒体帧分片总长度与frame_size不一致");
        }
        remove(key, frame);
        return new CompletedFrame(frame.firstHeader, frame.payload);
    }

    public int expire(long nowMs) {
        int expired = 0;
        Iterator<Map.Entry<FrameKey, PartialFrame>> iterator = frames.entrySet().iterator();
        while (iterator.hasNext()) {
            PartialFrame frame = iterator.next().getValue();
            if (nowMs - frame.updatedAtMs >= timeoutMs) {
                allocatedBytes -= frame.payload.length;
                iterator.remove();
                expired++;
            }
        }
        return expired;
    }

    public int clear() {
        int size = frames.size();
        frames.clear();
        allocatedBytes = 0;
        return size;
    }

    public int inFlightFrames() { return frames.size(); }
    public long allocatedBytes() { return allocatedBytes; }

    private void validate(DatagramHeader header) {
        if (sessionEpoch == 0 || header.getSessionEpoch() != sessionEpoch) {
            throw new IllegalArgumentException("媒体DATAGRAM不属于当前媒体会话");
        }
        if (header.getFrameSize() > maximumFrameBytes) {
            throw new IllegalArgumentException("媒体帧超过单帧大小上限");
        }
    }

    private void reserve(int frameBytes) {
        while (!frames.isEmpty()
                && (frames.size() >= maximumFrames || allocatedBytes + frameBytes > maximumBytes)) {
            Map.Entry<FrameKey, PartialFrame> oldest = frames.entrySet().iterator().next();
            remove(oldest.getKey(), oldest.getValue());
        }
        if (frames.size() >= maximumFrames || allocatedBytes + frameBytes > maximumBytes) {
            throw new IllegalStateException("媒体重组缓冲区容量不足");
        }
    }

    private void remove(FrameKey key, PartialFrame frame) {
        if (frames.remove(key) != null) {
            allocatedBytes -= frame.payload.length;
        }
    }

    private static boolean overlaps(PartialFrame frame, int offset, int length) {
        long end = (long) offset + length;
        for (int index = frame.received.nextSetBit(0); index >= 0;
             index = frame.received.nextSetBit(index + 1)) {
            long existingEnd = (long) frame.offsets[index] + frame.lengths[index];
            if ((long) offset < existingEnd && (long) frame.offsets[index] < end) {
                return true;
            }
        }
        return false;
    }

    private static boolean equalsBytes(ByteBuf source, int sourceIndex,
                                       byte[] target, int targetIndex, int length) {
        for (int i = 0; i < length; i++) {
            if (source.getByte(sourceIndex + i) != target[targetIndex + i]) {
                return false;
            }
        }
        return true;
    }

    private static final class FrameKey {
        private final long trackId;
        private final long frameSequence;

        private FrameKey(long trackId, long frameSequence) {
            this.trackId = trackId;
            this.frameSequence = frameSequence;
        }

        @Override
        public boolean equals(Object other) {
            if (this == other) return true;
            if (!(other instanceof FrameKey)) return false;
            FrameKey key = (FrameKey) other;
            return trackId == key.trackId && frameSequence == key.frameSequence;
        }

        @Override
        public int hashCode() {
            int result = (int) (trackId ^ (trackId >>> 32));
            result = 31 * result + (int) (frameSequence ^ (frameSequence >>> 32));
            return result;
        }
    }

    private static final class PartialFrame {
        private final DatagramHeader firstHeader;
        private final byte[] payload;
        private final BitSet received;
        private final int[] offsets;
        private final int[] lengths;
        private int receivedBytes;
        private long updatedAtMs;

        private PartialFrame(DatagramHeader header, long nowMs) {
            firstHeader = header;
            payload = new byte[header.getFrameSize()];
            received = new BitSet(header.getFragmentCount());
            offsets = new int[header.getFragmentCount()];
            lengths = new int[header.getFragmentCount()];
            updatedAtMs = nowMs;
        }

        private boolean matches(DatagramHeader header) {
            return firstHeader.getSessionEpoch() == header.getSessionEpoch()
                    && firstHeader.getTrackId() == header.getTrackId()
                    && firstHeader.getFrameSequence() == header.getFrameSequence()
                    && firstHeader.getFrameSize() == header.getFrameSize()
                    && firstHeader.getFragmentCount() == header.getFragmentCount()
                    && firstHeader.getKind() == header.getKind()
                    && firstHeader.getFlags() == header.getFlags()
                    && firstHeader.getCodecGenerationToken() == header.getCodecGenerationToken()
                    && firstHeader.getCaptureTimestampUs() == header.getCaptureTimestampUs();
        }
    }

    public static final class CompletedFrame {
        private final DatagramHeader header;
        private final byte[] payload;

        private CompletedFrame(DatagramHeader header, byte[] payload) {
            this.header = header;
            this.payload = payload;
        }

        public DatagramHeader getHeader() { return header; }
        public byte[] getPayload() { return payload; }
    }
}
