import logging
import time
from queue import Queue
import traceback
import threading
import numpy as np

import cv2
from mindx.sdk import base
from mindx.sdk.base import VideoDecoder, VideoDecodeConfig, VdecCallBacker

import faulthandler
# 在import之后直接添加以下启用代码即可
faulthandler.enable()
# 后边正常写你的代码

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

DEVICE_ID = 0
SKIP_INTERVAL = 1
VIDEO_HEIGHT = 1080
VIDEO_WIDTH = 1920
VIDEO_BREAK_FRAME = 100
# MODEL_IMG_FORMAT = base.nv12
VIDEO_PATH = "fireDetection.264"
# VIDEO_PATH = "rtsp://128.8.84.224:18554/station/stream1"


def process0_video_decode(decode_q_pro: Queue, video_stream_path: str, channelId: int):
    """
    decode_q_pro: Queue, decoded frames to be put in
    video_stream_path: str, video to decode from
    channelId: int, video decode channel
    """
    try:
        logger.info(
            "=========================== Start process0_video_decode ===========================")
        time.sleep(0.1)
        def vdec_callback(decodeImage, channelId, frameId):
            """
            Define videoDecoder callback function
            """
            logger.info(
                f"process0_video_decode decodeImage channelId: {channelId}, frameId: {frameId} width: {decodeImage.width}, height: {decodeImage.height}")
            decodeImage.to_host()
            logger.info(
                f"process0_video_decode to be put in queue resized_image: {decodeImage.width} * {decodeImage.height}")
            decode_q_pro.put(
                (frameId, decodeImage, decodeImage.original_height, decodeImage.original_width))
            logger.info("process0_video_decode producer queue, elements {}".format(
                decode_q_pro.qsize()))
        # 1. initialize video decoder callback
        vdec_callbacker = VdecCallBacker()
        vdec_callbacker.registerVdecCallBack(vdec_callback)
        # 2. initialize video decoder config
        vdec_config = VideoDecodeConfig()
        vdec_config.skipInterval = SKIP_INTERVAL
        # NHWC (If use aipp should be base.nv12)
        vdec_config.inputVideoFormat = base.h264_main_level
        vdec_config.outputImageFormat = base.nv12
        vdec_config.width = VIDEO_WIDTH
        vdec_config.height = VIDEO_HEIGHT
        # 3. initialize video decode object
        video_decoder = VideoDecoder(
            vdec_config, vdec_callbacker, deviceId=DEVICE_ID, channelId=channelId)
        logger.info(
            f"=========================== Start 1 process0_video_decode {video_stream_path} ===========================")
        # 4. use ffmpeg av and VideoDecoder to decode video
        import av
        with av.open(video_stream_path) as container:
            logger.info(
                f"============================process0_video_decode start decode stream =========================")
            for frame_id, packet in enumerate(container.demux()):
                logger.info(
                    f"process0_video_decode curr frame id: {frame_id}, size: {packet.size}")
                if packet.size == 0:
                    break
                try:
                    logger.info(
                        f"process0_video_decode start decode ")
                    video_decoder.decode(packet, frame_id)
                    logger.info(f"==== process0_video_decode decode success")
                except Exception as err:
                    logger.error("process0_video_decode decode error: ", err)
                if frame_id == VIDEO_BREAK_FRAME:
                    break
        logger.info(
            "=============== Finish process0_video_decode ========================")
        del video_decoder
        # 避免快速创建吃满显存
        time.sleep(5)
    except Exception as e:
        print("process0_video_decode failed.", repr(e))
        traceback.print_exc()


def main_pool():
    NUM_DECODE_PROCESS = 1
    q_decode = Queue()
    threads = []
    for i in range(NUM_DECODE_PROCESS):
        t1 = threading.Thread(target=process0_video_decode, args=(q_decode, VIDEO_PATH, i,))
        threads.append(t1)
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    print(f"==== start get image ====")
    print(f"==== q_decode len is {q_decode.qsize()} ====")
    idx=0
    while 1:
        if q_decode.empty():
            break
        else:
            image = q_decode.get()
            print(f"type image is {type(image[1])}, {image[1]}")
            im = np.array(image[1].to_tensor())
            print(f"image shape is {im.shape}")
            # print(f"image is {im}")
            cv2.imwrite(f"image_dir/{idx}.jpg",im[0])
        idx += 1

if __name__ == '__main__':
    base.mx_init()
    main_pool()