from traceback import format_exc

import cv2
import numpy as np
from loguru import logger
import subprocess as sp
from common import Constant
from enums.ExceptionEnum import ExceptionType
from exception.CustomerException import ServiceException
from util.Cv2Utils import push_video_stream, clear_pull_p

pull_p = None
push_p =None
p_push_status=[0,0]

def start_pull_p(pull_url, requestId):
    try:
        command = ['ffmpeg']
        # if pull_url.startswith("rtsp://"):
        #     command.extend(['-timeout', '20000000', '-rtsp_transport', 'tcp'])
        # if pull_url.startswith("http") or pull_url.startswith("rtmp"):
        #     command.extend(['-rw_timeout', '20000000'])
        command.extend(['-re',
                        '-y',
                        '-an',
                        # '-hwaccel', 'cuda', cuvid
                        '-c:v', 'h264_cuvid',
                        # '-resize', self.wah,
                        '-i', pull_url,
                        '-f', 'rawvideo',
                        # '-pix_fmt', 'bgr24',
                        '-r', '25',
                        '-'])
        return sp.Popen(command, stdout=sp.PIPE)
    except ServiceException as s:
        logger.error("构建拉流管道异常: {}, requestId:{}", s.msg, requestId)
        raise s
    except Exception as e:
        logger.error("构建拉流管道异常：{}, requestId:{}", format_exc(), requestId)
        raise e

def pull_read_video_stream(pull_p, pull_url, width, height, width_height_3, w_2, h_2, requestId):
    result = None
    try:
        if pull_p is None:
            pull_p = start_pull_p(pull_url, requestId)
        in_bytes = pull_p.stdout.read(width_height_3)
        if in_bytes is not None and len(in_bytes) > 0:
            try:
                # result = (np.frombuffer(in_bytes, np.uint8).reshape([height * 3 // 2, width, 3]))
                # ValueError: cannot reshape array of size 3110400 into shape (1080,1920)
                result = (np.frombuffer(in_bytes, np.uint8)).reshape((height, width))
                result = cv2.cvtColor(result, cv2.COLOR_YUV2BGR_NV12)
                # result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
                if result.shape[1] > Constant.width:
                    result = cv2.resize(result, (result.shape[1] // 2, result.shape[0] // 2), interpolation=cv2.INTER_LINEAR)
            except Exception:
                logger.error("视频格式异常：{}, requestId:{}", format_exc(), requestId)
                raise ServiceException(ExceptionType.VIDEO_RESOLUTION_EXCEPTION.value[0],
                                       ExceptionType.VIDEO_RESOLUTION_EXCEPTION.value[1])
    except ServiceException as s:
        clear_pull_p(pull_p, requestId)
        raise s
    except Exception:
        clear_pull_p(pull_p, requestId)
        pull_p = None
        width = None
        height = None
        width_height_3 = None
        logger.error("读流异常：{}, requestId:{}", format_exc(), requestId)
    return result, pull_p, width, height, width_height_3

while True:
    frame, pull_p, width, height, width_height_3 = pull_read_video_stream(pull_p, 'rtmp://live.play.t-aaron.com/live/THSAr', 1920,
                                                                          1080*3//2, 1920*1080*3//2, 960, 540,
                                                                          '111')
    if frame is not None:
        push_p = push_video_stream(frame, push_p, 'rtmp://live.push.t-aaron.com/live/THSAs', p_push_status, '11')
clear_pull_p(pull_p, "11111")
close_all_p(push_p, None, None, "11111")