import logging
import time
from queue import Queue
import traceback
import threading

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


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

DEVICE_ID = 0
SKIP_INTERVAL = 1
VIDEO_HEIGHT = 1536
VIDEO_WIDTH = 2048
VIDEO_BREAK_FRAME = 100
MODEL_IMG_FORMAT = base.nv12
VIDEO_PATH = "./test1.mp4"


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.outputImageFormat = MODEL_IMG_FORMAT
        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(
            "=========================== Start 1 process0_video_decode ===========================")

        # 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):
        # if USE_IMGS:
        #     t = threading.Thread(target=process0_read_image, args=(image_processor, q_resize, i,))
        #     threads.append(t)
        # else:
        #     t1 = threading.Thread(target=process0_video_decode, args=(q_decode, VIDEO_PATH, i,))
        #     threads.append(t1)
        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()

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