guzijian / wvp-GB28181-pro · Files

package com.genersoft.iot.vmp.jtt1078.publisher;

import com.genersoft.iot.vmp.jtt1078.codec.AudioCodec;
import com.genersoft.iot.vmp.jtt1078.entity.Media;
import com.genersoft.iot.vmp.jtt1078.entity.MediaEncoding;
import com.genersoft.iot.vmp.jtt1078.flv.FlvEncoder;
import com.genersoft.iot.vmp.jtt1078.subscriber.RTMPPublisher;
import com.genersoft.iot.vmp.jtt1078.subscriber.Subscriber;
import com.genersoft.iot.vmp.jtt1078.subscriber.VideoSubscriber;
import com.genersoft.iot.vmp.jtt1078.util.ByteHolder;
import com.genersoft.iot.vmp.jtt1078.util.Configs;
import io.netty.channel.ChannelHandlerContext;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;
import java.util.Iterator;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

public class Channel
{
    static Logger logger = LoggerFactory.getLogger(Channel.class);

    ConcurrentLinkedQueue<Subscriber> subscribers;

    // [恢复] FFmpeg 推流进程管理器
    RTMPPublisher rtmpPublisher;

    String tag;
    boolean publishing;
    ByteHolder buffer;
    AudioCodec audioCodec;
    FlvEncoder flvEncoder;
    private long firstTimestamp = -1;

    // ========== 新增：视频参数检测和FFmpeg自动重启相关 ==========

    /** 上一次接收到的SPS哈希值，用于检测视频参数变化 */
    private int lastSPSHash = 0;

    /** FFmpeg是否需要重启的标志 */
    private AtomicBoolean ffmpegNeedRestart = new AtomicBoolean(false);

    /** 最后一次FFmpeg重启时间（毫秒），用于防止频繁重启 */
    private AtomicLong lastRestartTime = new AtomicLong(0);

    /** 重启冷却时间（毫秒），3秒内不重复重启 */
    private static final long RESTART_COOLDOWN_MS = 3000;

    /** 是否正在等待新的I帧（FFmpeg重启期间） */
    private AtomicBoolean waitingForIFrame = new AtomicBoolean(false);

    /** 视频参数信息 */
    private volatile VideoParamInfo currentVideoParam = new VideoParamInfo();

    /** 视频参数内部类 */
    private static class VideoParamInfo {
        int width = 0;
        int height = 0;
        int fps = 0;
        String resolution = "unknown";

        @Override
        public String toString() {
            return String.format("%dx%d@%dfps [%s]", width, height, fps, resolution);
        }
    }

    public Channel(String tag)
    {
        this.tag = tag;
        this.subscribers = new ConcurrentLinkedQueue<>();
        this.flvEncoder = new FlvEncoder(true, true);
        this.buffer = new ByteHolder(2048 * 100);

        // [恢复] 启动 FFmpeg 进程
        // 只要配置了 rtmp.url，就启动 FFmpeg 去拉取当前的 HTTP 流并转码推送
        String rtmpUrl = Configs.get("rtmp.url");
        if (StringUtils.isNotBlank(rtmpUrl))
        {
            logger.info("[{}] 启动 FFmpeg 进程推流至: {}", tag, rtmpUrl);
            rtmpPublisher = new RTMPPublisher(tag);
            rtmpPublisher.start();
        }
    }

    public boolean isPublishing()
    {
        return publishing;
    }

    public Subscriber subscribe(ChannelHandlerContext ctx)
    {
        Subscriber subscriber = new VideoSubscriber(this.tag, ctx);
        this.subscribers.add(subscriber);
        return subscriber;
    }

    public void writeAudio(long timestamp, int pt, byte[] data)
    {
        // 1. 转码为 PCM (用于 FLV 封装，FFmpeg 会从 FLV 中读取 PCM 并转码为 AAC)
        if (audioCodec == null)
        {
            audioCodec = AudioCodec.getCodec(pt);
            logger.info("[{}] audio codec: {}", tag, MediaEncoding.getEncoding(Media.Type.Audio, pt));
        }
        // 写入到内部广播，FFmpeg 通过 HTTP 拉取这个数据
        broadcastAudio(timestamp, audioCodec.toPCM(data));
    }

    public void writeVideo(long sequence, long timeoffset, int payloadType, byte[] h264)
    {
        if (firstTimestamp == -1) firstTimestamp = timeoffset;
        this.publishing = true;
        this.buffer.write(h264);

        while (true)
        {
            byte[] nalu = readNalu();
            if (nalu == null) break;
            if (nalu.length < 4) continue;

            // ========== 新增：检测视频参数变化 ==========
            int nalType = nalu[4] & 0x1F;
            boolean isIDR = (nalType == 5);  // IDR帧（I帧）

            if (nalType == 7) {  // SPS (Sequence Parameter Set)
                checkAndHandleSPSChange(nalu);
            }

            // 如果正在等待I帧，跳过所有数据直到收到新的I帧
            if (waitingForIFrame.get()) {
                if (isIDR) {
                    logger.info("[{}] 收到新的I帧，停止等待，FFmpeg应已重启完成", tag);
                    waitingForIFrame.set(false);
                    // 重置FLV编码器
                    this.flvEncoder = new FlvEncoder(true, true);
                    firstTimestamp = timeoffset;
                } else {
                    // 跳过非I帧数据
                    continue;
                }
            }

            // 1. 封装为 FLV Tag (必须)
            // FFmpeg 通过 HTTP 读取这些 FLV Tag
            byte[] flvTag = this.flvEncoder.write(nalu, (int) (timeoffset - firstTimestamp));

            if (flvTag == null) continue;

            // 广播给所有的观众
            broadcastVideo(timeoffset, flvTag);
        }
    }

    /**
     * 检查SPS变化并处理FFmpeg重启
     */
    private synchronized void checkAndHandleSPSChange(byte[] nalu) {
        int currentHash = Arrays.hashCode(nalu);

        if (lastSPSHash != 0 && currentHash != lastSPSHash) {
            // SPS变化了，说明视频参数发生了变化
            // 注意：即使解析出来的分辨率/帧率相同，SPS的其它变化（如profile、level、VUI参数等）也会导致解码问题
            // 因此：只要SPS哈希变化，就重启FFmpeg

            // 解析新的视频参数（仅用于日志显示）
            VideoParamInfo newParam = parseVideoParam(nalu);
            VideoParamInfo oldParam = currentVideoParam;

            logger.info("[{}] ====== 检测到SPS变化，需要重启FFmpeg ======", tag);
            logger.info("[{}] 旧参数: {}", tag, oldParam);
            logger.info("[{}] 新参数: {}", tag, newParam);
            logger.info("[{}] SPS哈希变化: {} -> {}", tag,
                    Integer.toHexString(lastSPSHash),
                    Integer.toHexString(currentHash));

            // 检查冷却时间
            long now = System.currentTimeMillis();
            long timeSinceLastRestart = now - lastRestartTime.get();

            if (timeSinceLastRestart < RESTART_COOLDOWN_MS) {
                logger.warn("[{}] FFmpeg重启过于频繁（距上次{}ms），跳过本次重启",
                        tag, timeSinceLastRestart);
                // 即使跳过重启，仍需更新SPS哈希
                lastSPSHash = currentHash;
                currentVideoParam = newParam;
                return;
            }

            logger.info("[{}] ====== 准备重启FFmpeg ======", tag);
            logger.info("[{}] 原因: SPS参数变化（可能包含分辨率、帧率、profile、level等）", tag);
            logger.info("[{}] 预计中断时间: 1-2秒", tag);

            // 标记需要重启
            ffmpegNeedRestart.set(true);

            // 立即触发重启（异步执行，避免阻塞）
            final String currentTag = tag;
            Thread restartThread = new Thread(() -> {
                try {
                    logger.info("[{}] SPS变化检测到，立即触发FFmpeg重启...", currentTag);
                    restartRtmpPublisher();
                } catch (Exception e) {
                    logger.error("[{}] 触发FFmpeg重启失败: {}", currentTag, e.getMessage(), e);
                }
            }, "FFmpeg-Restart-" + tag);
            restartThread.setDaemon(true);
            restartThread.start();
        }

        // 更新SPS哈希和视频参数
        lastSPSHash = currentHash;
        currentVideoParam = parseVideoParam(nalu);
    }

    /**
     * 解析H.264 SPS获取视频参数
     */
    private VideoParamInfo parseVideoParam(byte[] sps) {
        VideoParamInfo info = new VideoParamInfo();
        try {
            // H.264 SPS解析
            // SPS格式: [NAL头(1字节) + profile_idc(1) + constraints(1) + level_idc(1) + 5字节对齐 + sps数据]

            if (sps.length < 6) {
                logger.warn("[{}] SPS数据长度不足，无法解析: {}", tag, sps.length);
                return info;
            }

            int profile_idc = sps[4] & 0xFF;
            int constraint_flags = sps[5] & 0xFF;
            int level_idc = sps[6] & 0xFF;

            // 查找SPS结束位置（开始于67或68）
            int spsStart = -1;
            for (int i = 4; i < Math.min(sps.length - 4, 10); i++) {
                if ((sps[i] & 0x1F) == 7) {
                    spsStart = i;
                    break;
                }
            }

            if (spsStart == -1) {
                // 简化模式：直接使用SPS长度估算
                int spsLen = (sps.length > 10) ? 4 : 1;
                for (int i = 4; i < sps.length - 1; i++) {
                    if (sps[i] == 0 && sps[i+1] == 0 && sps[i+2] == 0 && sps[i+3] == 1) {
                        spsLen = i - 4;
                        break;
                    }
                }

                // 根据SPS长度估算分辨率（非常粗略）
                int lenCategory = spsLen / 10;
                if (spsLen < 20) {
                    info.width = 352;
                    info.height = 288;
                    info.fps = 15;
                } else if (spsLen < 40) {
                    info.width = 704;
                    info.height = 576;
                    info.fps = 25;
                } else {
                    info.width = 1280;
                    info.height = 720;
                    info.fps = 25;
                }
            } else {
                // 详细解析bitstream
                // 这里简化处理，实际项目中可以使用完整的H.264解析库
                info.width = 1280;  // 默认值
                info.height = 720;
                info.fps = 25;
            }

            // 根据SPS长度特征判断分辨率
            // 这是一个简化的估算方法
            int estimatedSize = sps.length;

            // 子码流特征: CIF (352x288) -> SPS约8-15字节
            // 主码流特征: 720P (1280x720) -> SPS约16-30字节
            // 1080P (1920x1080) -> SPS约30+字节
            if (estimatedSize < 20) {
                // CIF 或类似分辨率
                if (estimatedSize < 12) {
                    info.width = 352;
                    info.height = 288;
                    info.fps = 15;
                    info.resolution = "CIF";
                } else {
                    info.width = 704;
                    info.height = 576;
                    info.fps = 25;
                    info.resolution = "4CIF";
                }
            } else if (estimatedSize < 35) {
                // 720P 或类似
                info.width = 1280;
                info.height = 720;
                info.fps = 25;
                info.resolution = "720P";
            } else {
                // 1080P 或更高
                info.width = 1920;
                info.height = 1080;
                info.fps = 30;
                info.resolution = "1080P";
            }

        } catch (Exception e) {
            logger.error("[{}] 解析SPS失败: {}", tag, e.getMessage());
        }

        return info;
    }

    /**
     * 重启FFmpeg推流进程
     */
    public void restartRtmpPublisher() {
        long now = System.currentTimeMillis();
        lastRestartTime.set(now);

        logger.info("[{}] ====== 开始重启FFmpeg推流 ======", tag);
        logger.info("[{}] 当前时间: {}", tag, now);

        try {
            // 1. 标记等待I帧
            waitingForIFrame.set(true);

            // 2. 关闭旧的FFmpeg进程
            logger.info("[{}] 步骤1: 关闭旧FFmpeg进程...", tag);
            if (rtmpPublisher != null) {
                try {
                    rtmpPublisher.close();
                } catch (Exception e) {
                    logger.warn("[{}] 关闭旧FFmpeg进程时出错: {}", tag, e.getMessage());
                }
                rtmpPublisher = null;
            }

            // 3. 等待FFmpeg进程完全关闭
            logger.info("[{}] 步骤2: 等待FFmpeg进程关闭...", tag);
            Thread.sleep(500);

            // 4. 清空缓冲区，准备接收新的视频数据
            logger.info("[{}] 步骤3: 清空视频缓冲区...", tag);
            buffer.clear();
            firstTimestamp = -1;

            // 5. 重置FLV编码器
            logger.info("[{}] 步骤4: 重置FLV编码器...", tag);
            flvEncoder = new FlvEncoder(true, true);

            // 6. 重新启动FFmpeg进程
            logger.info("[{}] 步骤5: 启动新FFmpeg进程...", tag);
            String rtmpUrl = Configs.get("rtmp.url");
            if (StringUtils.isNotBlank(rtmpUrl)) {
                rtmpPublisher = new RTMPPublisher(tag);
                rtmpPublisher.start();
                logger.info("[{}] 新FFmpeg进程已启动", tag);
            } else {
                logger.warn("[{}] 未配置rtmp.url，跳过启动", tag);
            }

            // 7. 重置标志
            ffmpegNeedRestart.set(false);

            logger.info("[{}] ====== FFmpeg重启完成 ======", tag);
            logger.info("[{}] 请等待1-2秒让FFmpeg完成初始化...", tag);

        } catch (Exception e) {
            logger.error("[{}] 重启FFmpeg失败: {}", tag, e.getMessage(), e);
            waitingForIFrame.set(false);
            ffmpegNeedRestart.set(false);
            // 确保rtmpPublisher被清空
            rtmpPublisher = null;
        }
    }

    /**
     * 外部调用：主动触发码流切换（对应1078的9102指令）
     * 调用此方法后，FFmpeg会在检测到视频参数变化时自动重启
     */
    public void notifyStreamSwitch() {
        logger.info("[{}] ====== 收到码流切换通知 ======", tag);
        logger.info("[{}] 即将切换码流，请等待设备响应...", tag);

        // 记录切换前的状态
        VideoParamInfo beforeSwitch = currentVideoParam;
        logger.info("[{}] 切换前视频参数: {}", tag, beforeSwitch);

        // 重置SPS哈希，这样一旦设备发来新的SPS就能立即检测到
        lastSPSHash = 0;

        // 标记需要重启
        ffmpegNeedRestart.set(true);

        // 启动一个线程来处理后续的重启逻辑
        new Thread(() -> {
            try {
                // 等待设备切换并发送新的I帧
                Thread.sleep(2000);

                // 如果还没有收到新的I帧，手动触发一次检查
                if (waitingForIFrame.get()) {
                    logger.warn("[{}] 未在2秒内收到新I帧，检查是否需要手动重启...", tag);
                    // 这里不直接重启，而是等待下一个SPS自动触发
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }, "StreamSwitch-Watcher-" + tag).start();
    }

    /**
     * 获取当前视频参数（用于监控和调试）
     */
    public VideoParamInfo getCurrentVideoParam() {
        return currentVideoParam;
    }

    /**
     * 获取FFmpeg是否需要重启
     */
    public boolean isFFmpegNeedRestart() {
        return ffmpegNeedRestart.get();
    }

    public void broadcastVideo(long timeoffset, byte[] flvTag)
    {
        for (Subscriber subscriber : subscribers)
        {
            subscriber.onVideoData(timeoffset, flvTag, flvEncoder);
        }
    }

    public void broadcastAudio(long timeoffset, byte[] flvTag)
    {
        for (Subscriber subscriber : subscribers)
        {
            subscriber.onAudioData(timeoffset, flvTag, flvEncoder);
        }
    }

    public void unsubscribe(long watcherId)
    {
        for (Iterator<Subscriber> itr = subscribers.iterator(); itr.hasNext(); )
        {
            Subscriber subscriber = itr.next();
            if (subscriber.getId() == watcherId)
            {
                itr.remove();
                subscriber.close();
                return;
            }
        }
    }

    public void close()
    {
        logger.info("[{}] 关闭Channel，开始清理资源...", tag);

        for (Iterator<Subscriber> itr = subscribers.iterator(); itr.hasNext(); )
        {
            Subscriber subscriber = itr.next();
            subscriber.close();
            itr.remove();
        }

        // 关闭 FFmpeg 进程
        if (rtmpPublisher != null) {
            logger.info("[{}] 关闭FFmpeg推流进程...", tag);
            rtmpPublisher.close();
            rtmpPublisher = null;
        }

        logger.info("[{}] Channel已关闭", tag);
    }

    // [恢复] 原版 readNalu (FFmpeg 偏好带 StartCode 的数据，或者 FlvEncoder 需要)
    private byte[] readNalu()
    {
        // 寻找 00 00 00 01
        for (int i = 0; i < buffer.size() - 3; i++)
        {
            int a = buffer.get(i + 0) & 0xff;
            int b = buffer.get(i + 1) & 0xff;
            int c = buffer.get(i + 2) & 0xff;
            int d = buffer.get(i + 3) & 0xff;
            if (a == 0x00 && b == 0x00 && c == 0x00 && d == 0x01)
            {
                if (i == 0) continue;
                byte[] nalu = new byte[i];
                buffer.sliceInto(nalu, i);
                return nalu;
            }
        }
        return null;
    }
}