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

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.List;

/**
 * GB28181 标准 PS 封装（与 gb28181_streampackageForH264 一致，FFmpeg -f mpeg 可解）。
 * <p>
 * I 帧：Pack + System + PSM + PES(Annex-B)<br>
 * P 帧：Pack + PES(Annex-B)<br>
 * 超大帧按 65522 字节分片，每片独立 Pack（I 帧每片带 System+PSM）。
 */
public final class Gb28181PsMuxer {

    /** gb28181 PS_PES_PAYLOAD_SIZE */
    private static final int MAX_PES_PAYLOAD = 65522;
    private static final int VIDEO_STREAM_ID = 0xE0;

    private Gb28181PsMuxer() {
    }

    public static byte[] mux(List<byte[]> nals, long pts90k, boolean keyFrame) throws IOException {
        if (nals == null || nals.isEmpty()) {
            return new byte[0];
        }
        return muxAnnexB(toAnnexB(nals), pts90k, keyFrame);
    }

    static byte[] muxAnnexB(byte[] annexB, long pts90k, boolean keyFrame) throws IOException {
        if (annexB == null || annexB.length == 0) {
            return new byte[0];
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream(annexB.length + 256);
        int offset = 0;
        while (offset < annexB.length) {
            int chunk = Math.min(annexB.length - offset, MAX_PES_PAYLOAD);
            writePsChunk(out, annexB, offset, chunk, pts90k, keyFrame);
            offset += chunk;
        }
        return out.toByteArray();
    }

    private static void writePsChunk(ByteArrayOutputStream out, byte[] annexB, int offset, int length,
                                     long pts90k, boolean keyFrame) throws IOException {
        out.write(Gb28181PsHeaders.buildStandardPackHeader(pts90k));
        if (keyFrame) {
            out.write(Gb28181PsHeaders.buildSystemHeader());
            out.write(Gb28181PsHeaders.buildStandardPsm());
        }
        out.write(Gb28181PsHeaders.buildPesHeader(VIDEO_STREAM_ID, length, pts90k));
        out.write(annexB, offset, length);
    }

    private static byte[] toAnnexB(List<byte[]> nals) throws IOException {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        for (byte[] nal : nals) {
            if (nal != null && nal.length > 0) {
                out.write(new byte[]{0x00, 0x00, 0x00, 0x01});
                out.write(nal);
            }
        }
        return out.toByteArray();
    }
}
