import time

import acl
import cv2
import subprocess as sp
import sys
import threading

import numpy as np

# enum
memcpy_kind = {
    "ACL_MEMCPY_HOST_TO_HOST": 0,
    "ACL_MEMCPY_HOST_TO_DEVICE": 1,
    "ACL_MEMCPY_DEVICE_TO_HOST": 2,
    "ACL_MEMCPY_DEVICE_TO_DEVICE": 3
}

"""
enType 0
0 H265 main level
1 H264 baseline level
2 H264 main level
3 H264 high level
"""
VENC_ENTYPE = 2
"""
format 2
1 YUV420 semi-planner (nv12)
2 YVU420 semi-planner (nv21)
"""
VENC_FORMAT = 1

DEVICE_ID = 0
IMAGE_INPUT_WIDTH = 1920
IMAGE_INPUT_HEIGHT = 1080

MJPEG_URL = "http://172.24.12.243:33378"
PUSH_URL = "rtsp://127.0.0.1:8554/video"

cmd = ["ffmpeg",
       "-re",
       "-max_delay", "100",
       "-stream_loop", "-1",
       "-i", "-",
       "-c", "copy",
       "-level:v", "1",
       # "-preset", "ultrafast",
       "-b:v", "400k", "-bufsize", "4000k", "-maxrate", "8000k",
       "-f", "rtsp",
       "-an",
       # "-v","debug",
       "-timeout", "10",  # 设置TCP连接的等待时间
       "-rtsp_transport", "tcp",  # 使用TCP推流
       PUSH_URL]  # rtsp port 8554


def t_g(frame):
    # 获取图像的高、宽
    height, width, _ = frame.shape

    # 将 BGR 图像转换为 YUV 格式
    yuv = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV)

    # 创建一个大小为 height + height//2 和 width 的空数组来存储 YUV420SP (NV12) 格式
    yuv420sp = np.zeros((height + height // 2, width), dtype=np.uint8)

    # Y 分量：直接将 Y 分量拷贝到 yuv420sp 中
    yuv420sp[:height, :] = yuv[:, :, 0]

    # UV 分量（NV12格式，U和V交替排列）
    u = yuv[:, :, 1]
    v = yuv[:, :, 2]

    # 将 U 和 V 分量合并为一个交替存储的格式（NV12：U和V交替）
    yuv420sp[height:, ::2] = u[::2, ::2]  # U 分量
    yuv420sp[height:, 1::2] = v[::2, ::2]  # V 分量
    return yuv420sp

def check_ret(message, ret_int):
    """
    功能简介：检测pyACL函数返回值是否正常，如果非0则会抛出异常
    参数：ret，pyACL函数返回值
    返回值：无
    """
    if ret_int != 0:
        raise Exception("{} failed ret_int={}"
                        .format(message, ret_int))


def check_none(message, ret_none):
    """
    功能简介：检测值是否为空，如果为空则会抛出异常
    参数：ret_none
    返回值：无
    """
    if ret_none is None:
        raise Exception("{} failed"
                        .format(message))


data_list = []


class AclVenc(object):
    """
    视频编码
    """

    def __init__(self):
        self.input_stream_mem = None
        self.venc_channel_desc = 0
        self.dvpp_pic_desc = 0
        self.frame_config = None
        self.cb_thread_id = None
        self.pic_host = None
        self.context, ret = acl.rt.create_context(DEVICE_ID)
        check_ret("acl.rt.create_context", ret)
        self.callback_run_flag = True
        self.stream_host = None
        self.spPipe = sp.Popen(cmd, shell=False, stdin=sp.PIPE, stderr=sp.PIPE)
        self.lock = threading.Lock()
        self.cont = 0

    def release_resource(self):
        if self.stream_host:
            acl.rt.free_host(self.stream_host)
        if self.pic_host:
            acl.rt.free_host(self.pic_host)
        if self.input_stream_mem:
            acl.media.dvpp_free(self.input_stream_mem)
        if self.venc_channel_desc != 0:
            acl.media.venc_destroy_channel_desc(self.venc_channel_desc)
        if self.dvpp_pic_desc != 0:
            acl.media.dvpp_destroy_pic_desc(self.dvpp_pic_desc)
        if self.frame_config:
            acl.media.venc_destroy_frame_config(self.frame_config)
        acl.rt.destroy_context(self.context)
        self.spPipe.stdin.close()
        self.spPipe.wait()

    def venc_init(self):
        self.venc_channel_desc = acl.media.venc_create_channel_desc()
        check_none("acl.media.venc_create_channel_desc",
                   self.venc_channel_desc)

    def cb_thread_func(self, args_list):
        context = args_list[0]
        timeout = args_list[1]
        print("[INFO] cb_thread_func args_list = ", context, timeout,
              self.callback_run_flag)

        context, ret = acl.rt.create_context(DEVICE_ID)
        if ret != 0:
            print("[INFO] cb_thread_func acl.rt.create_context ret=", ret)
            return
        while self.callback_run_flag is True:
            ret = acl.rt.process_report(timeout)

        ret = acl.rt.destroy_context(context)
        print("[INFO] cb_thread_func acl.rt.destroy_context ret=", ret)

    def callback_func(self, input_pic_desc, output_stream_desc, user_data):
        print("===== callback start =======")
        self.lock.acquire()
        if output_stream_desc == 0:
            print("[INFO] [venc] output_stream_desc is null")
            return
        stream_data = acl.media.dvpp_get_stream_desc_data(output_stream_desc)
        if stream_data is None:
            print("[INFO] [venc] acl.media.dvpp_get_stream_desc_data is none")
            return
        ret = acl.media.dvpp_get_stream_desc_ret_code(output_stream_desc)
        if ret == 0:
            stream_data_size = acl.media.dvpp_get_stream_desc_size(
                output_stream_desc)
            print("[INFO] [venc] stream_data size", stream_data_size)
            # stream memcpy d2h
            data = bytes(stream_data_size)
            data_ptr = acl.util.bytes_to_ptr(data)
            ret = acl.rt.memcpy(data_ptr, stream_data_size,
                                stream_data, stream_data_size,
                                memcpy_kind.get("ACL_MEMCPY_DEVICE_TO_HOST"))
            if ret != 0:
                print("[INFO] [venc] acl.rt.memcpy ret=", ret)
                return
            data_list.append(data)

            if self.spPipe != None:
                try:
                    self.spPipe.stdin.write(data_list[0])
                    print("Push num {} done!!!".format(str(self.cont)))
                    del data_list[0]
                    self.cont += 1
                except Exception as e:
                    print("[ERROR]", e)
            else:
                print("pipe break *******")
        else:
            print("[ERROR] [venc] ret", ret)
        self.lock.release()
        print("==== callback end =====")

    def venc_set_desc(self, width, height):
        # venc_channel_desc set function
        acl.media.venc_set_channel_desc_thread_id(
            self.venc_channel_desc, self.cb_thread_id)

        acl.media.venc_set_channel_desc_callback(
            self.venc_channel_desc, self.callback_func)
        acl.media.venc_set_channel_desc_entype(
            self.venc_channel_desc, VENC_ENTYPE)
        acl.media.venc_set_channel_desc_pic_format(
            self.venc_channel_desc, VENC_FORMAT)
        key_frame_interval = 16
        acl.media.venc_set_channel_desc_key_frame_interval(
            self.venc_channel_desc, key_frame_interval)
        acl.media.venc_set_channel_desc_pic_height(
            self.venc_channel_desc, height)
        acl.media.venc_set_channel_desc_pic_width(
            self.venc_channel_desc, width)

    def venc_set_frame_config(self, frame_confg, eos, iframe):
        acl.media.venc_set_frame_config_eos(frame_confg, eos)
        acl.media.venc_set_frame_config_force_i_frame(frame_confg, iframe)

    def venc_get_frame_config(self, frame_confg):
        get_eos = acl.media.venc_get_frame_config_eos(frame_confg)
        check_ret("acl.media.venc_get_frame_config_eos", get_eos)
        get_force_frame = acl.media.venc_get_frame_config_force_i_frame(
            frame_confg)
        check_ret("acl.media.venc_get_frame_config_force_i_frame",
                  get_force_frame)

    def venc_run(self):
        timeout = 1000
        self.cb_thread_id, ret = acl.util.start_thread(
            self.cb_thread_func, [self.context, timeout])
        check_ret("acl.util.start_thread", ret)
        print("[INFO] start_thread", self.cb_thread_id, ret)
        self.venc_set_desc(IMAGE_INPUT_WIDTH, IMAGE_INPUT_HEIGHT)
        print("[INFO] set venc channel desc")

        ret = acl.media.venc_create_channel(self.venc_channel_desc)
        check_ret("acl.media.venc_create_channel", ret)

        # 初始化摄像头
        self.capture = cv2.VideoCapture(0)
        if not self.capture.isOpened():
            raise Exception("Failed to open camera")

        # 设置摄像头分辨率
        self.capture.set(cv2.CAP_PROP_FRAME_WIDTH, IMAGE_INPUT_WIDTH)
        self.capture.set(cv2.CAP_PROP_FRAME_HEIGHT, IMAGE_INPUT_HEIGHT)

        self.capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc(*"MJPG"))  # 设置宽度
        self.capture.set(cv2.CAP_PROP_FPS, 120)
        self.capture.set(cv2.CAP_PROP_AUTO_WB, 1.0)  # 打开自动白平衡
        for i in range(40):
           self.capture.read()
           time.sleep(1000 / 25 / 1000)  # 目的是为了清掉相机里白平衡的数据?
        self.capture.set(cv2.CAP_PROP_GAIN, 255)
        self.capture.set(cv2.CAP_PROP_AUTO_EXPOSURE, 3.0)  # 打开自动曝光 (相机会记住曝光)
        # self.capture.set(cv2.CAP_PROP_AUTO_EXPOSURE, 1.0)  # 关闭自动曝光 (相机会记住曝光)
        # self.capture.set(cv2.CAP_PROP_EXPOSURE, 10000 * 2 ** 7)
        print("[INFO] Camera initialized")
        #acl.rt.set_op_execute_time_out(50) 无法使用
        while True:
            ret, frame = self.capture.read()
            # resized_image = dvpp.jpegd(image_input)
            yuv_image = t_g(frame)
            #yuv_image = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV_I420)
            # resized_image = dvpp.resize(yuv_image, IMAGE_INPUT_WIDTH, IMAGE_INPUT_HEIGHT)
            file_size = yuv_image.size
            file_context = yuv_image

            if "bytes_to_ptr" in dir(acl.util):
                bytes_data = file_context.tobytes()
                file_mem = acl.util.bytes_to_ptr(bytes_data)
            else:
                file_mem = acl.util.numpy_to_ptr(file_context)

            input_size = file_size
            input_mem, ret = acl.media.dvpp_malloc(input_size)
            check_ret("acl.media.dvpp_malloc", ret)

            ret = acl.rt.memcpy(
                input_mem, input_size, file_mem,
                file_size, memcpy_kind.get("ACL_MEMCPY_HOST_TO_DEVICE"))
            check_ret("acl.rt.memcpy", ret)

            self.frame_config = acl.media.venc_create_frame_config()
            check_none("acl.media.venc_create_frame_config", self.frame_config)
            print("[INFO] create_frame_config")

            self.dvpp_pic_desc = acl.media.dvpp_create_pic_desc()
            check_none("acl.media.dvpp_create_pic_desc", self.dvpp_pic_desc)
            ret = acl.media.dvpp_set_pic_desc_data(self.dvpp_pic_desc, input_mem)
            ret = acl.media.dvpp_set_pic_desc_size(self.dvpp_pic_desc, input_size)
            print("[INFO] set pic desc size")

            self.venc_set_frame_config(self.frame_config, 0, 0)
            print("[INFO] set frame config")
            self.venc_get_frame_config(self.frame_config)

            ret = acl.media.venc_send_frame(
                    self.venc_channel_desc,
                    self.dvpp_pic_desc, 0, self.frame_config, None)

            check_ret("acl.media.venc_send_frame", ret)

            print("[INFO] venc send frame success")

            ret = acl.media.dvpp_free(input_mem)
            check_ret("acl.media.dvpp_free", ret)

        print("[INFO] venc send frame eos")
        self.venc_set_frame_config(self.frame_config, 1, 0)
        ret = acl.media.venc_send_frame(
            self.venc_channel_desc, 0, 0, self.frame_config, None)
        print("[INFO] acl.media.venc_send_frame ret=", ret)

        ret = acl.media.venc_destroy_channel(self.venc_channel_desc)
        check_ret("acl.media.venc_destroy_channel", ret)

        self.thread_join()
        print("[INFO] thread join")

    def venc_stream_desc_set(self):
        stream_format = 2
        timestamp = 123456
        ret_code = 1
        eos = 1

        # stream desc
        dvpp_stream_desc = acl.media.dvpp_create_stream_desc()
        check_none("acl.media.dvpp_create_stream_desc", dvpp_stream_desc)

        # stream_desc set function
        acl.media.dvpp_set_stream_desc_format(dvpp_stream_desc, stream_format)
        acl.media.dvpp_set_stream_desc_timestamp(dvpp_stream_desc, timestamp)
        acl.media.dvpp_set_stream_desc_ret_code(dvpp_stream_desc, ret_code)
        acl.media.dvpp_set_stream_desc_eos(dvpp_stream_desc, eos)

        ret = acl.media.dvpp_destroy_stream_desc(dvpp_stream_desc)
        check_ret("acl.media.dvpp_destroy_stream_desc", ret)

    def thread_join(self):
        self.callback_run_flag = False
        ret = acl.util.stop_thread(self.cb_thread_id)
        print("[INFO] stop_thread", ret)


def args_check():
    args = sys.argv
    if len(args) == 3:
        global MJPEG_URL
        MJPEG_URL = args[1]
        global PUSH_URL
        PUSH_URL = args[2]
    else:
        print("use default params")


if __name__ == '__main__':
    args_check()
    ret = acl.init("")
    check_ret("acl.init", ret)
    ret = acl.rt.set_device(DEVICE_ID)
    check_ret("acl.rt.set_device", ret)
    run_mode, ret = acl.rt.get_run_mode()
    check_ret("acl.rt.get_run_mode", ret)
    venc_handle = AclVenc()
    venc_handle.venc_init()
    venc_handle.venc_run()
    venc_handle.venc_stream_desc_set()
    venc_handle.release_resource()
    ret = acl.rt.reset_device(DEVICE_ID)
    check_ret("acl.rt.reset_device", ret)
    ret = acl.finalize()
    check_ret("acl.finalize", ret)