import asyncio
import base64
import traceback
from asyncio import Event
from copy import deepcopy, copy
from datetime import datetime

import cv2

from algo.model_manager import AlgoModelExec
from algo.stream_loader import OpenCVStreamLoad
from common.device_status_manager import DeviceStatusManager
from common.global_logger import logger
from common.global_thread_pool import GlobalThreadPool
from common.http_utils import send_request
from common.image_plotting import Annotator
from entity.device import Device

from ultralytics import YOLO

from scene_handler.base_scene_handler import BaseSceneHandler
from services.global_config import GlobalConfig
from tcp.tcp_manager import TcpManager

COLOR_RED = (0, 0, 255)
COLOR_GREEN = (255, 0, 0)

ALARM_DICT = {
    'hat_and_mask': {
        'alarmType': '11',
        'alarmContent': '未佩戴呼吸防护设备',
        'alarmSoundMessage': b'\xaa\x01\x00\x93\x12\x00\xA6',
        'label': '未佩戴呼吸防护设备'
    },
    'no_jiandu': {
        'alarmType': '12',
        'alarmContent': '没有监护人员',
        'alarmSoundMessage': b'\xaa\x01\x00\x93\x13\x00\xA7',
        'label': '没有监护人员'
    },
    'break': {
        'alarmType': '3',
        'alarmContent': '非法闯入',
        'alarmSoundMessage': b'\xaa\x01\x00\x93\x00\x00\x94',
        'label': '非法闯入'
    },
    'smoke': {
        'alarmType': '6',
        'alarmContent': '吸烟',
        'alarmSoundMessage': b'\xaa\x01\x00\x93\x03\x00\x97',
        'label': '吸烟'
    },
    'no_blower': {
        'alarmType': '13',
        'alarmContent': '没有检测到通风设备',
        'alarmSoundMessage': b'\xaa\x01\x00\x93\x1A\x00\xAE',
        'label': '没有检测到通风设备'
    },
    'no_extinguisher': {
        'alarmType': '14',
        'alarmContent': '没有检测到灭火器',
        'alarmSoundMessage': b'\xaa\x01\x00\x93\x1B\x00\xAF',
        'label': '没有检测到灭火器'
    }
}


def intersection_area(bbox1, bbox2):
    # 计算两个坐标框的重叠面积
    x1, y1, x2, y2 = bbox1
    x3, y3, x4, y4 = bbox2

    xi1 = max(x1, x3)
    yi1 = max(y1, y3)
    xi2 = min(x2, x4)
    yi2 = min(y2, y4)

    width = max(0, xi2 - xi1)
    height = max(0, yi2 - yi1)

    return width * height


def bbox_area(bbox):
    # 计算坐标框的面积
    x1, y1, x2, y2 = bbox
    return (x2 - x1) * (y2 - y1)


def get_person_head(person_bbox, heads):
    best_head = None
    max_overlap = 0

    for head in heads:
        head_bbox = head.xyxy.cpu().squeeze()
        overlap_area = intersection_area(person_bbox, head_bbox)
        head_area = bbox_area(head_bbox)
        overlap_ratio = overlap_area / head_area

        if overlap_ratio >= 0.8 and (best_head is None or overlap_area > max_overlap or (
                overlap_area == max_overlap and float(head.conf) > float(best_head.conf))):
            best_head = head
            max_overlap = overlap_area
    return best_head


def is_overlapping(bbox1, bbox2):
    # 检查两个坐标框是否重叠
    x1, y1, x2, y2 = bbox1
    x3, y3, x4, y4 = bbox2

    return not (x2 < x3 or x4 < x1 or y2 < y3 or y4 < y1)


def image_to_base64(numpy_image, format='jpg'):
    # 将NumPy数组转换为图片格式的字节流
    # format指定图片格式，'png'或'jpeg'
    success, encoded_image = cv2.imencode(f'.{format}', numpy_image)
    if not success:
        raise ValueError("Could not encode image")

    # 将字节流转换为Base64编码
    base64_encoded_image = base64.b64encode(encoded_image)

    # 将bytes类型转换为UTF-8字符串
    base64_message = base64_encoded_image.decode('utf-8')

    return base64_message


def handle_alarm_info(type, frame, frame_alarm, person_box, conf=None, person_id=None):
    if type not in frame_alarm or frame_alarm[type] is None:
        annotator = Annotator(deepcopy(frame), None, 18, "Arial.ttf", False, example="人")
        frame_alarm[type] = {'count': 0, 'annotator': annotator}
    frame_alarm[type]['count'] = frame_alarm[type]['count'] + 1
    alarm_annotator = frame_alarm[type]['annotator']
    if person_box is not None:
        alarm_annotator.box_label(person_box,
                                  ALARM_DICT[type]['label'] + (f'{conf:.2f}' if conf is not None else '') + (
                                      f' id={person_id}' if person_id is not None else ''),
                                  color=COLOR_RED,
                                  rotated=False)
    return alarm_annotator


class LimitSpaceSceneHandler(BaseSceneHandler):

    def __init__(self, device: Device, thread_id: str, tcp_manager: TcpManager, main_loop, range_points):
        super().__init__(device=device, thread_id=thread_id, tcp_manager=tcp_manager, main_loop=main_loop)
        # self.device = device
        # self.thread_id = thread_id
        self.stream_loader = OpenCVStreamLoad(camera_url=device.input_stream_url, camera_code=device.code,
                                              device_thread_id=thread_id)
        self.device_status_manager = DeviceStatusManager()
        self.thread_pool = GlobalThreadPool()

        self.model = YOLO('weights/labor-v8-20241114.pt')
        self.model_classes = {
            0: '三脚架',
            3: '人',
            4: '作业信息公示牌',
            6: '危险告知牌',
            9: '反光衣',
            11: '呼吸面罩',
            13: '四合一',
            15: '头',
            16: '安全告知牌',
            18: '安全帽',
            20: '安全标识牌',
            24: '工服',
            34: '灭火器',
            43: '警戒线',
            58: '鼓风机',
        }

        self.alarm_interval_dict = {}
        self.alarm_interval = device.alarm_interval

        self.socket_interval_dict = {}
        self.socket_interval = device.alarm_interval
        self.socket_retry = 3

        self.__stop_event = Event(loop=main_loop)  # 使用 Event 控制线程的运行状态

    def stop_task(self, **kwargs):
        logger.info(f'stop detection task {self.device.id}, thread_id: {self.thread_id}')
        self.__stop_event.set()
        self.stream_loader.stop()  # 停止视频流加载的线程

    def send_tcp_message(self, message: bytes, have_response=False):
        asyncio.run_coroutine_threadsafe(
            self.tcp_manager.send_message_to_device(device_id=self.device.id,
                                                    message=message,
                                                    have_response=have_response),
            self.main_loop)

    def send_alarm_message(self, type):
        if self.tcp_manager:
            if self.socket_interval_dict.get(type) is None \
                    or (datetime.now() - self.socket_interval_dict.get(type)).total_seconds() > int(self.socket_interval):
                logger.debug("send alarm message %s %s", ALARM_DICT[type]['alarmContent'],
                             ALARM_DICT[type]['alarmSoundMessage'])
                self.send_tcp_message(ALARM_DICT[type]['alarmSoundMessage'], have_response=True)
                self.socket_interval_dict[type] = datetime.now()

    def send_alarm_record(self, type, frame_alarm):
        if self.alarm_interval < 0:
            return

        global_config = GlobalConfig()
        push_config = global_config.get_alarm_push_config()
        if push_config and push_config.push_url:
            if self.alarm_interval_dict.get(type) is None \
                    or (datetime.now() - self.alarm_interval_dict.get(type)).total_seconds() > int(self.alarm_interval):
                logger.debug("send alarm record")

                annotator_result = frame_alarm[type]['annotator'].result()
                alarm_image = deepcopy(annotator_result)

                data = {}
                data["device_id"] = self.device.id
                data["alarm_type"] = ALARM_DICT[type]['alarmType']
                data["alarm_content"] = ALARM_DICT[type]['alarmContent']
                data["alarm_value"] = frame_alarm[type]['count']
                data["alarm_image"] = image_to_base64(alarm_image)
                data["alarm_time"] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

                url = push_config.push_url
                send_request(url, data)

                data_copy = copy(data)
                # 从拷贝的字典中移除"alarm_image"键
                data_copy.pop("alarm_image", None)
                logger.debug(f"send to {url}: {data_copy}")

                self.alarm_interval_dict[type] = datetime.now()

    def model_predict(self, frame):
        results_gen = self.model(frame, save_txt=False, save=False, verbose=False, conf=0.5,
                                 classes=list(self.model_classes.keys()),
                                 imgsz=640,
                                 stream=True)
        results = list(results_gen)  # 确保生成器转换为列表
        result = results[0]
        result_boxes = [box for box in result.boxes]
        pred_ids = [int(box.cls) for box in result_boxes]
        pred_names = [self.model_classes[int(box.cls)] for box in result_boxes]
        return result_boxes, pred_ids, pred_names

    def process_alarm(self, frame, result_boxes, pred_ids, pred_names, frame_alarm):
        persons = [box for box in result_boxes if int(box.cls) == 3]
        helmets = [box for box in result_boxes if int(box.cls) == 18]
        heads = [box for box in result_boxes if int(box.cls) == 15]

        has_jianduyuan = False
        has_others = False

        for person in persons:
            person_bbox = person.xyxy.cpu().squeeze()

            # 检查这个人是否佩戴了安全帽
            has_helmet = True
            person_head = get_person_head(person_bbox, heads)
            if person_head is not None:
                has_helmet = any(
                    is_overlapping(person_head.xyxy.cpu().squeeze(), helmet.xyxy.cpu().squeeze()) for helmet in
                    helmets)

            if not has_helmet:
                has_others = True
                alarm_annotator = handle_alarm_info('break', frame, frame_alarm, person_bbox,
                                                         float(person.conf))
            else:
                has_jianduyuan = True

        if not has_jianduyuan:
            alarm_annotator = handle_alarm_info('no_jiandu', frame, frame_alarm, None)
            self.send_alarm_message('no_jiandu')
        if has_others:
            self.send_alarm_message('break')

    def process_labor(self, frame, result_boxes, pred_ids, pred_names):
        pass

    def run(self):
        while not self.stream_loader.init:
            if self.__stop_event.is_set():
                break  # 如果触发了停止事件，则退出循环
            self.stream_loader.init_cap()
        for frames in self.stream_loader:
            try:
                if self.__stop_event.is_set():
                    break  # 如果触发了停止事件，则退出循环
                # print('frame')
                if frames is None:
                    continue

                self.device_status_manager.set_status(device_id=self.device.id)
                # result_boxes, pred_ids, pred_names = self.model_predict(frames)

                frame_alarm = {}
                # self.process_alarm(frame, result_boxes, pred_ids, pred_names, frame_alarm)
                # self.process_labor(frame, result_boxes, pred_ids, pred_names)

                if len(frame_alarm.keys()) > 0:
                    for key in frame_alarm.keys():
                        if frame_alarm[key]['count'] > 0:
                            self.thread_pool.submit_task(self.send_alarm_record, key, frame_alarm, )

            except Exception as ex:
                traceback.print_exc()
                logger.error(ex)
