# -*- coding: utf-8 -*-
from concurrent.futures import ThreadPoolExecutor
from os.path import join
from threading import Thread
from traceback import format_exc

import cv2
import numpy as np
from loguru import logger
from util.Cv2Utils import write_or_video, write_ai_video, push_video_stream, close_all_p, video_conjuncing
from util.ImageUtils import url2Array, add_water_pic
from util.PlotsUtils import draw_painting_joint
from util.QueUtil import put_queue


class OnPushStreamThread(Thread):
    __slots__ = ('_msg', '_push_queue', '_context', 'ex', '_logo', '_image_queue')

    def __init__(self, *args):
        super().__init__()
        # 传参
        self._msg, self._push_queue, self._image_queue, self._context = args
        # 自带参数
        self.ex = None
        self._logo = None
        if self._context["video"]["video_add_water"]:
            self._logo = self._msg.get("logo_url")
            if self._logo:
                self._logo = url2Array(self._logo, enable_ex=False)
            if not self._logo:
                self._logo = cv2.imread(join(self._context['base_dir'], "image/logo.png"), -1)

    def run(self):
        request_id, push_queue, image_queue = self._msg.get("request_id"), self._push_queue, self._image_queue
        orFilePath, aiFilePath, logo = self._context.get("orFilePath"), self._context.get("aiFilePath"), self._logo
        or_video_file, ai_video_file, push_p = None, None, None
        push_url = self._msg.get("push_url")
        try:
            logger.info("开始启动推流线程！requestId:{}", request_id)
            with ThreadPoolExecutor(max_workers=2) as t:
                p_push_status, or_write_status, ai_write_status = [0, 0], [0, 0], [0, 0]
                while True:
                    push_parm = push_queue.get()
                    if push_parm is not None:
                        # [(1, 原视频帧, 分析视频帧)]
                        # # [视频帧、当前帧数、 总帧数、 [(问题数组、code、allowedList、label_arraylist、rainbows)]]
                        # res = (1, (pull_frame[1], pull_frame[2], pull_frame[3], []))
                        # [(2, 操作指令)]
                        if push_parm[0] == 1:  # 视频帧操作
                            frame, current_frame, all_frames, ques_list = push_parm[1]
                            copy_frame = frame.copy()
                            det_xywh = {}
                            if len(ques_list) > 0:
                                for qs in ques_list:
                                    det_xywh[qs[1]] = {}
                                    detect_targets_code = int(qs[0][0])
                                    score = qs[0][-1]
                                    label_array = qs[3][detect_targets_code]
                                    color = qs[4][detect_targets_code]
                                    if not isinstance(qs[0][1], (list, tuple, np.ndarray)):
                                        xc, yc, x2, y2 = int(qs[0][1]), int(qs[0][2]), int(qs[0][3]), int(qs[0][4])
                                        box = [(xc, yc), (x2, yc), (x2, y2), (xc, y2)]
                                    else:
                                        box = qs[0][1]
                                    draw_painting_joint(box, copy_frame, label_array, score, color, "leftTop")
                                    cd = det_xywh[qs[1]].get(detect_targets_code)
                                    if cd is None:
                                        det_xywh[qs[1]][detect_targets_code] = [
                                            [detect_targets_code, box, score, label_array, color]]
                                    else:
                                        det_xywh[qs[1]][detect_targets_code].append(
                                            [detect_targets_code, box, score, label_array, color])
                            if logo:
                                frame = add_water_pic(frame, logo, request_id)
                                copy_frame = add_water_pic(copy_frame, logo, request_id)
                            frame_merge = video_conjuncing(frame, copy_frame)
                            # 写原视频到本地
                            write_or_video_result = t.submit(write_or_video, frame, orFilePath, or_video_file,
                                                             or_write_status, request_id)
                            # 写识别视频到本地
                            write_ai_video_result = t.submit(write_ai_video, frame_merge, aiFilePath, ai_video_file,
                                                             ai_write_status, request_id)
                            if len(det_xywh) > 0:
                                put_queue(image_queue, (1, (det_xywh, frame, current_frame, all_frames)))
                            push_p = push_video_stream(frame_merge, push_p, push_url, p_push_status, request_id)
                            ai_video_file = write_ai_video_result.result()
                            or_video_file = write_or_video_result.result()
                        if push_parm[0] == 2:
                            if 'stop' == push_parm[1]:
                                logger.info("停止推流线程, requestId: {}", request_id)
                                close_all_p(push_p, or_video_file, ai_video_file, request_id)
                                or_video_file, ai_video_file, push_p = None, None, None
                                break
        except Exception as e:
            logger.error("推流线程异常:{}, requestId:{}", format_exc(), request_id)
            self.ex = e
        finally:
            close_all_p(push_p, or_video_file, ai_video_file, request_id)
        logger.info("推流线程停止完成！requestId:{}", request_id)


class OffPushStreamThread(Thread):
    __slots__ = ('_msg', '_push_queue', '_context', 'ex', '_logo', '_image_queue')

    def __init__(self, *args):
        super().__init__()
        # 传参
        self._msg, self._push_queue, self._image_queue, self._context = args
        # 自带参数
        self.ex = None
        self._logo = None
        if self._context["video"]["video_add_water"]:
            self._logo = self._msg.get("logo_url")
            if self._logo:
                self._logo = url2Array(self._logo, enable_ex=False)
            if not self._logo:
                self._logo = cv2.imread(join(self._context['base_dir'], "image/logo.png"), -1)

    def run(self):
        request_id, push_queue, image_queue = self._msg.get("request_id"), self._push_queue, self._image_queue
        aiFilePath, logo = self._context.get("aiFilePath"), self._logo
        ai_video_file, push_p = None, None
        push_url = self._msg.get("push_url")
        try:
            logger.info("开始启动推流线程！requestId:{}", request_id)
            with ThreadPoolExecutor(max_workers=1) as t:
                p_push_status, or_write_status, ai_write_status = [0, 0], [0, 0], [0, 0]
                while True:
                    push_parm = push_queue.get()
                    if push_parm is not None:
                        # [(1, 原视频帧, 分析视频帧)]
                        # # [视频帧、当前帧数、 总帧数、 [(问题数组、code、allowedList、label_arraylist、rainbows)]]
                        # res = (1, (pull_frame[1], pull_frame[2], pull_frame[3], []))
                        # [(2, 操作指令)]
                        if push_parm[0] == 1:  # 视频帧操作
                            frame, current_frame, all_frames, ques_list = push_parm[1]
                            copy_frame = frame.copy()
                            det_xywh = {}
                            if len(ques_list) > 0:
                                for qs in ques_list:
                                    det_xywh[qs[1]] = {}
                                    detect_targets_code = int(qs[0][0])
                                    score = qs[0][-1]
                                    label_array = qs[3][detect_targets_code]
                                    color = qs[4][detect_targets_code]
                                    if not isinstance(qs[0][1], (list, tuple, np.ndarray)):
                                        xc, yc, x2, y2 = int(qs[0][1]), int(qs[0][2]), int(qs[0][3]), int(qs[0][4])
                                        box = [(xc, yc), (x2, yc), (x2, y2), (xc, y2)]
                                    else:
                                        box = qs[0][1]
                                    draw_painting_joint(box, copy_frame, label_array, score, color, "leftTop")
                                    cd = det_xywh[qs[1]].get(detect_targets_code)
                                    if cd is None:
                                        det_xywh[qs[1]][detect_targets_code] = [
                                            [detect_targets_code, box, score, label_array, color]]
                                    else:
                                        det_xywh[qs[1]][detect_targets_code].append(
                                            [detect_targets_code, box, score, label_array, color])
                            if logo:
                                frame = add_water_pic(frame, logo, request_id)
                                copy_frame = add_water_pic(copy_frame, logo, request_id)
                            frame_merge = video_conjuncing(frame, copy_frame)
                            # 写识别视频到本地
                            write_ai_video_result = t.submit(write_ai_video, frame_merge, aiFilePath, ai_video_file,
                                                             ai_write_status, request_id)
                            if len(det_xywh) > 0:
                                put_queue(image_queue, (1, (det_xywh, frame, current_frame, all_frames)))
                            push_p = push_video_stream(frame_merge, push_p, push_url, p_push_status, request_id)
                            ai_video_file = write_ai_video_result.result()
                        if push_parm[0] == 2:
                            if 'stop' == push_parm[1]:
                                logger.info("停止推流线程, requestId: {}", request_id)
                                close_all_p(push_p, None, ai_video_file, request_id)
                                ai_video_file, push_p = None, None
                                break
        except Exception as e:
            logger.error("推流线程异常:{}, requestId:{}", format_exc(), request_id)
            self.ex = e
        finally:
            close_all_p(push_p, None, ai_video_file, request_id)
        logger.info("推流线程停止完成！requestId:{}", request_id)