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

import com.usky.cdi.service.sip.rtp.AnnexBParser;

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

/**
 * 组装 Annex-B H264 访问单元，保证 I 帧前 SPS/PPS 顺序正确。
 */
public final class H264AccessUnitBuilder {

    private byte[] sps;
    private byte[] pps;
    private boolean sentFirstKeyFrame;

    public void updateParameterSets(List<byte[]> nals) {
        for (byte[] nal : nals) {
            if (nal.length == 0) {
                continue;
            }
            int type = nal[0] & 0x1F;
            if (type == 7) {
                sps = nal;
            } else if (type == 8) {
                pps = nal;
            }
        }
    }

    public boolean hasParameterSets() {
        return sps != null && pps != null;
    }

    public void setSps(byte[] sps) {
        if (sps != null && sps.length > 0) {
            this.sps = sps;
        }
    }

    public void setPps(byte[] pps) {
        if (pps != null && pps.length > 0) {
            this.pps = pps;
        }
    }

    public boolean isParameterSetOnly(List<byte[]> nals) {
        if (nals.isEmpty()) {
            return true;
        }
        for (byte[] nal : nals) {
            if (nal.length == 0) {
                continue;
            }
            int type = nal[0] & 0x1F;
            if (type != 7 && type != 8) {
                return false;
            }
        }
        return true;
    }

    public boolean isKeyFrame(List<byte[]> nals) {
        for (byte[] nal : nals) {
            if (nal.length > 0 && (nal[0] & 0x1F) == 5) {
                return true;
            }
        }
        return false;
    }

    /**
     * @return Annex-B 访问单元；未就绪返回 null
     */
    public byte[] build(byte[] packetData, boolean keyFrameHint) throws IOException {
        List<byte[]> nals = AnnexBParser.extractNals(packetData);
        updateParameterSets(nals);
        if (nals.isEmpty() || isParameterSetOnly(nals)) {
            return null;
        }

        boolean keyFrame = keyFrameHint || isKeyFrame(nals);
        if (!keyFrame && !sentFirstKeyFrame) {
            return null;
        }
        if (keyFrame && !hasParameterSets() && !(containsType(nals, 7) && containsType(nals, 8))) {
            return null;
        }

        List<byte[]> ordered = orderNals(nals, keyFrame);
        sentFirstKeyFrame = sentFirstKeyFrame || keyFrame;
        return toAnnexB(ordered);
    }

    private List<byte[]> orderNals(List<byte[]> nals, boolean keyFrame) {
        List<byte[]> ordered = new ArrayList<>();
        if (keyFrame) {
            if (sps != null && !containsType(nals, 7)) {
                ordered.add(sps);
            }
            if (pps != null && !containsType(nals, 8)) {
                ordered.add(pps);
            }
        }
        List<byte[]> sorted = new ArrayList<>(nals);
        sorted.sort(Comparator.comparingInt(H264AccessUnitBuilder::nalPriority));
        for (byte[] nal : sorted) {
            if (nal.length == 0) {
                continue;
            }
            int type = nal[0] & 0x1F;
            if (type == 7 && ordered.stream().anyMatch(n -> (n[0] & 0x1F) == 7)) {
                continue;
            }
            if (type == 8 && ordered.stream().anyMatch(n -> (n[0] & 0x1F) == 8)) {
                continue;
            }
            ordered.add(nal);
        }
        return ordered;
    }

    private static int nalPriority(byte[] nal) {
        if (nal.length == 0) {
            return 99;
        }
        switch (nal[0] & 0x1F) {
            case 7:
                return 0;
            case 8:
                return 1;
            case 5:
                return 2;
            default:
                return 3;
        }
    }

    private static boolean containsType(List<byte[]> nals, int type) {
        for (byte[] nal : nals) {
            if (nal.length > 0 && (nal[0] & 0x1F) == type) {
                return true;
            }
        }
        return false;
    }

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