import json
import time
import subprocess as sp
import ffmpeg
import cv2
import sys
import random

import numpy as np

"""
获取视频基本信息
"""


def get_video_info(in_file):
    try:
        probe = ffmpeg.probe(in_file)
        # format = probe['format']
        # size = int(format['size'])/1024/1024
        video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
        if video_stream is None:
            print('No video stream found', file=sys.stderr)
            return
        width = int(video_stream['width'])
        height = int(video_stream['height'])
        # num_frames = int(video_stream['nb_frames'])
        up, down = str(video_stream['r_frame_rate']).split('/')
        video_stream['fps']=eval(up) / eval(down)
        # duration = float(video_stream['duration'])
        # bit_rate = int(video_stream['bit_rate'])/1000
        print('width: {}'.format(width))
        print('height: {}'.format(height))
        # print('num_frames: {}'.format(num_frames))
        # print('bit_rate: {}k'.format(bit_rate))
        # print('fps: {}'.format(fps))
        # print('size: {}MB'.format(size))
        # print('duration: {}'.format(duration))
        return video_stream



    except Exception as err:
        print("aaaaaaaaaaaaaaaaaaaa", err)
        return None


if __name__ == '__main__':
    file_path = "rtmp://live.play.t-aaron.com/live/THSAa"
    while True:
        video_info = get_video_info(file_path)
        print(json.dumps(video_info))
        if video_info is None:
            time.sleep(2)
            continue
        else:
            break
    command = ['ffmpeg',
               # '-re',
               '-y',
               '-c:v', 'h264_cuvid',
               '-i', file_path,
               '-pix_fmt', 'bgr24',
               '-f', 'rawvideo',
               '-an',
               '-']
    width = int(video_info['width'])
    height = int(video_info['height'])
    fps = int(video_info['fps'])
    p = sp.Popen(command, stdout=sp.PIPE)
    ai_video_file = cv2.VideoWriter(r"C:\Users\chenyukun\Desktop\shipin\aaaa.mp4", cv2.VideoWriter_fourcc(*'mp4v'), fps,
                                    (width, height))
    start1 = time.time()
    while True:
        start = time.time()
        in_bytes = p.stdout.read(int(width * height * 3))
        if not in_bytes:
            print(in_bytes)
            p.stdout.close()
            p.wait()
            break
        # img = (np.frombuffer(in_bytes, np.uint8)).reshape((int(height), int(width)))
        # bgr_img = cv2.cvtColor(img, cv2.COLOR_YUV2BGR_NV12)
        in_frame = (np.frombuffer(in_bytes, np.uint8).reshape([int(height), int(width), 3]))
        # print("拉流时间：", time.time() - start)
        # frame = cv2.resize(in_frame, (1280, 720))   # 改变图片尺寸
        # i += 1
        # print(round(time.time()-start, 5))
        #
        ai_video_file.write(in_frame)
        # if time.time() - start1 > 60:
        #     ai_video_file.release()
        #     p.stdout.close()
        #     p.wait()
        #     break
        # cv2.imshow('frame', in_frame)
        # cv2.waitKey(1)
        # time.sleep(1111)
    p.stdout.close()
    p.wait()
    p.kill()