package com.zhiqiyun.open.camera.rtsp;

import com.baomidou.mybatisplus.core.toolkit.IdWorker;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.ffmpeg.avcodec.AVCodecParameters;
import org.bytedeco.ffmpeg.avformat.AVFormatContext;
import org.bytedeco.ffmpeg.avformat.AVStream;
import org.bytedeco.ffmpeg.global.avcodec;
import org.bytedeco.ffmpeg.global.avutil;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.FFmpegFrameRecorder;
import org.bytedeco.javacv.FFmpegLogCallback;
import org.bytedeco.javacv.Frame;

/**
 * @author xin yang
 * @date 2022/10/13
 */
@Slf4j
public class RtspPushSrs {

//    //内网推流地址
//    public final static String INTRANET_SRS_PUSH_ADDRESS = "rtmp://116.8.106.156:7011/live/livestream/%s";
//
//    //外网拉流地址
//    public final static String SRS_PUSH_ADDRESS = "http://116.8.106.156:9800/live/livestream/%s.flv";

    //内网推流地址
    public final static String INTRANET_SRS_PUSH_ADDRESS = "rtmp://192.168.0.187:1935/live/livestream/%s";

    //外网拉流地址
//    public final static String SRS_PUSH_ADDRESS = "http://222.84.255.194:9800/live/livestream/%s.flv";
    public final static String SRS_PUSH_ADDRESS = "https://wlapi.zhlc.liucheng.gov.cn/live/livestream/%s.flv";

    /**
     * 读取指定的rtsp视频流，推送到SRS服务器
     *
     * @param rtspId 视频流ID
     * @throws Exception
     */
    public static void grabberAndPush(String rtspId) throws Exception {

        String rtspUrl = RtspTransfer.PATH_MAP.get(rtspId);

        // ffmepg日志级别
        avutil.av_log_set_level(avutil.AV_LOG_ERROR);
        FFmpegLogCallback.set();

        // 实例化帧抓取器对象，将文件路径传入
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(rtspUrl);
        long startTime = System.currentTimeMillis();
//        grabber.setFormat("hevc");
        // 设置读取的最大数据，单位字节
//        grabber.setOption("probesize", "10000");
        // 设置分析的最长时间，单位微秒
//        grabber.setOption("analyzeduration", "20000");
        log.info("开始初始化帧抓取器");
        grabber.setOption("rtsp_transport", "tcp");
//        grabber.setOption("stimeout", "20000000");
        // 初始化帧抓取器，例如数据结构（时间戳、编码器上下文、帧对象等），
        // 如果入参等于true，还会调用avformat_find_stream_info方法获取流的信息，放入AVFormatContext类型的成员变量oc中
//        grabber.setImageWidth(640);
//        grabber.setImageHeight(480);
        try {
            grabber.start();
        } catch (Exception e) {
            log.error("grabberAndPush error：" + e.getMessage(), e);
        }
        log.info("帧抓取器初始化完成，耗时[{}]毫秒", System.currentTimeMillis() - startTime);
        // grabber.start方法中，初始化的解码器信息存在放在grabber的成员变量oc中
        AVFormatContext avFormatContext = grabber.getFormatContext();

        // 文件内有几个媒体流（一般是视频流+音频流）
        int streamNum = avFormatContext.nb_streams();

        // 没有媒体流就不用继续了
        if (streamNum < 1) {
            log.error("文件内不存在媒体流");
            return;
        }

        // 取得视频的帧率
        int frameRate = (int) grabber.getVideoFrameRate();

        log.info("视频帧率[{}]，视频时长[{}]秒，媒体流数量[{}]", frameRate, avFormatContext.duration() / 1000000, avFormatContext.nb_streams());

        // 遍历每一个流，检查其类型
        for (int i = 0; i < streamNum; i++) {
            AVStream avStream = avFormatContext.streams(i);
            AVCodecParameters avCodecParameters = avStream.codecpar();
            log.info("流的索引[{}]，编码器类型[{}]，编码器ID[{}]", i, avCodecParameters.codec_type(), avCodecParameters.codec_id());
        }

        // 视频宽度
        int frameWidth = grabber.getImageWidth();
        // 视频高度
        int frameHeight = grabber.getImageHeight();
        // 音频通道数量
        int audioChannels = grabber.getAudioChannels();
        //编码格式
//        int videoCodec = grabber.getVideoCodec();
        //封装格式
        String format = grabber.getFormat();

        log.info("视频宽度[{}]，视频高度[{}]，音频通道数[{}]，封装格式{{}]", frameWidth, frameHeight, audioChannels, format);

        // 实例化FFmpegFrameRecorder，将SRS的推送地址传入
        String rtmpUrl = String.format(INTRANET_SRS_PUSH_ADDRESS, rtspId);
        log.info("流媒体服务地址：{}", rtmpUrl);
        FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(rtmpUrl, frameWidth, frameHeight, audioChannels);
//        recorder.setInterleaved(true);
//        recorder.setVideoOption("crf", "28");
        // https://trac.ffmpeg.org/wiki/StreamingGuide)
        recorder.setVideoOption("tune", "zerolatency");
        // ultrafast对CPU消耗最低
        recorder.setVideoOption("preset", "ultrafast");
        // Constant Rate Factor (see: https://trac.ffmpeg.org/wiki/Encode/H.264)
        recorder.setVideoOption("crf", "28");
        // 设置编码格式
        recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
        // 设置封装格式
        recorder.setFormat("flv");
        // 一秒内的帧数
        recorder.setFrameRate(frameRate);
//         两个关键帧之间的帧数
        recorder.setGopSize(frameRate * 2);
        // 设置音频通道数，与视频源的通道数相等
        recorder.setAudioChannels(audioChannels);
        // yuv420p
        recorder.setPixelFormat(0);
        startTime = System.currentTimeMillis();
        log.info("开始初始化帧抓取器");

        // 初始化帧录制器，例如数据结构（音频流、视频流指针，编码器），
        // 调用av_guess_format方法，确定视频输出时的封装方式，
        // 媒体上下文对象的内存分配，
        // 编码器的各项参数设置
        recorder.start();

        log.info("帧录制初始化完成，耗时[{}]毫秒", System.currentTimeMillis() - startTime);

        Frame frame;

        startTime = System.currentTimeMillis();

        log.info("开始推流");

        long videoTS = 0;

        int videoFrameNum = 0;
        int audioFrameNum = 0;
        int dataFrameNum = 0;

        // 假设一秒钟15帧，那么两帧间隔就是(1000/15)毫秒
        int interVal = 1000 / frameRate;
        // 发送完一帧后sleep的时间，不能完全等于(1000/frameRate)，不然会卡顿，
        // 要更小一些，这里取八分之一
        interVal /= 8;

        // 持续从视频源取帧
        while (null != (frame = grabber.grab())) {
            if (videoTS == 0) {
                RtspTransfer.FRAME_GRABBER_CONCURRENT_HASH_MAP.put(rtspId, grabber);
                RtspTransfer.FRAME_RECORDER_CONCURRENT_HASH_MAP.put(rtspId, recorder);
            }

            videoTS = 1000 * (System.currentTimeMillis() - startTime);
            log.info("持续从视频源取帧 videoTS：{}", videoTS);

            // 时间戳
            recorder.setTimestamp(videoTS);

            // 有图像，就把视频帧加一
            if (null != frame.image) {
                videoFrameNum++;
            }

            // 有声音，就把音频帧加一
            if (null != frame.samples) {
                audioFrameNum++;
            }

            // 有数据，就把数据帧加一
            if (null != frame.data) {
                dataFrameNum++;
            }

            // 取出的每一帧，都推送到SRS
            recorder.record(frame);

            // 停顿一下再推送
            Thread.sleep(interVal);
        }

        log.info("推送完成，视频帧[{}]，音频帧[{}]，数据帧[{}]，耗时[{}]秒",
                videoFrameNum,
                audioFrameNum,
                dataFrameNum,
                (System.currentTimeMillis() - startTime) / 1000);

        // 关闭帧录制器
        recorder.close();
        // 关闭帧抓取器
        grabber.close();
    }

    public static void main(String[] args) throws Exception {
        String rtspId = IdWorker.get32UUID();
        System.out.println(rtspId);
        RtspTransfer.PATH_MAP.put(rtspId, "rtsp://wowzaec2demo.streamlock.net/vod/mp4:BigBuckBunny_115k.mp4");
        // rtmp://121.40.185.233:1935/live/livestream/9ee705f6d05676b02d036088c2f09718
        grabberAndPush(rtspId);
    }
}