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

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 从 FFmpeg AVPacket 数据中解析 H264 NAL（Annex B 或 AVCC）
 */
public final class AnnexBParser {

    private AnnexBParser() {
    }

    public static List<byte[]> extractNals(byte[] data) {
        if (data == null || data.length == 0) {
            return new ArrayList<>();
        }
        if (isAnnexB(data)) {
            return splitAnnexB(data);
        }
        return splitAvcc(data);
    }

    private static boolean isAnnexB(byte[] data) {
        return data.length >= 3 && data[0] == 0 && data[1] == 0
                && (data[2] == 1 || (data.length >= 4 && data[2] == 0 && data[3] == 1));
    }

    private static List<byte[]> splitAnnexB(byte[] data) {
        List<byte[]> nals = new ArrayList<>();
        List<Integer> codeStarts = new ArrayList<>();
        List<Integer> nalBegins = new ArrayList<>();
        for (int i = 0; i < data.length - 3; i++) {
            if (data[i] == 0 && data[i + 1] == 0 && data[i + 2] == 1) {
                codeStarts.add(i);
                nalBegins.add(i + 3);
                i += 2;
            } else if (i + 3 < data.length && data[i] == 0 && data[i + 1] == 0
                    && data[i + 2] == 0 && data[i + 3] == 1) {
                codeStarts.add(i);
                nalBegins.add(i + 4);
                i += 3;
            }
        }
        if (nalBegins.isEmpty()) {
            nals.add(data);
            return nals;
        }
        for (int s = 0; s < nalBegins.size(); s++) {
            int begin = nalBegins.get(s);
            int end = (s + 1 < codeStarts.size()) ? codeStarts.get(s + 1) : data.length;
            if (begin < end) {
                nals.add(Arrays.copyOfRange(data, begin, end));
            }
        }
        return nals;
    }

    private static List<byte[]> splitAvcc(byte[] data) {
        List<byte[]> nals = new ArrayList<>();
        int pos = 0;
        while (pos + 4 <= data.length) {
            int len = ((data[pos] & 0xFF) << 24) | ((data[pos + 1] & 0xFF) << 16)
                    | ((data[pos + 2] & 0xFF) << 8) | (data[pos + 3] & 0xFF);
            pos += 4;
            if (len <= 0 || pos + len > data.length) {
                break;
            }
            nals.add(Arrays.copyOfRange(data, pos, pos + len));
            pos += len;
        }
        if (nals.isEmpty() && data.length > 0) {
            nals.add(data);
        }
        return nals;
    }
}
