package com.usky.cdi.service.sip.rtp;

import java.io.IOException;
import java.io.OutputStream;

/**
 * H264 RTP 解包，支持 Single NAL、STAP-A、FU-A
 */
public class H264Depacketizer {

    private static final byte[] NAL_START_CODE = {0x00, 0x00, 0x00, 0x01};

    private int fuExpectedType = -1;

    public void depacketize(RtpPacket packet, OutputStream out) throws IOException {
        if (packet == null || packet.getPayloadLength() <= 0) {
            return;
        }
        byte[] data = packet.getPayload();
        int offset = packet.getPayloadOffset();
        int length = packet.getPayloadLength();
        int nalType = data[offset] & 0x1F;

        if (nalType >= 1 && nalType <= 23) {
            writeNal(out, data, offset, length);
        } else if (nalType == 24) {
            depacketizeStapA(data, offset, length, out);
        } else if (nalType == 28) {
            depacketizeFuA(data, offset, length, out);
        }
    }

    private void depacketizeStapA(byte[] data, int offset, int length, OutputStream out)
            throws IOException {
        int pos = offset + 1;
        int end = offset + length;
        while (pos + 2 <= end) {
            int nalSize = ((data[pos] & 0xFF) << 8) | (data[pos + 1] & 0xFF);
            pos += 2;
            if (pos + nalSize > end) {
                break;
            }
            writeNal(out, data, pos, nalSize);
            pos += nalSize;
        }
    }

    private void depacketizeFuA(byte[] data, int offset, int length, OutputStream out)
            throws IOException {
        if (length < 2) {
            return;
        }
        int fuHeader = data[offset + 1] & 0xFF;
        boolean start = (fuHeader & 0x80) != 0;
        int nalType = fuHeader & 0x1F;

        if (start) {
            fuExpectedType = nalType;
            int nalHeader = (data[offset] & 0xE0) | nalType;
            out.write(NAL_START_CODE);
            out.write((byte) nalHeader);
        } else if (nalType != fuExpectedType) {
            return;
        }
        out.write(data, offset + 2, length - 2);
        if ((fuHeader & 0x40) != 0) {
            fuExpectedType = -1;
        }
    }

    private void writeNal(OutputStream out, byte[] data, int offset, int length)
            throws IOException {
        out.write(NAL_START_CODE);
        out.write(data, offset, length);
    }
}
