import time

from sqlalchemy.orm import Session
from typing import List
from fastapi import UploadFile
import subprocess
import torch
from app.util.yolov5.models.common import DetectMultiBackend
from app.util.yolov5.utils.torch_utils import select_device
from app.util.yolov5.utils.dataloaders import LoadStreams
from app.util.yolov5.utils.general import check_img_size, Profile, non_max_suppression, cv2, scale_boxes
from ultralytics.utils.plotting import Annotator, colors

from app.model.crud import project_detect_crud as pdc
from app.model.schemas.project_detect_schemas import ProjectDetectIn, ProjectDetectOut, ProjectDetectLogIn
from app.model.bussiness_model import ProjectDetect, ProjectDetectImg, ProjectTrain, ProjectDetectLog, ProjectDetectLogImg
from app.util.random_utils import random_str
from app.config.config_reader import detect_url
from app.util import os_utils as os
from app.util import random_utils as ru
from app.config.config_reader import yolo_url
from app.websocket.web_socket_server import room_manager


def add_detect(detect_in: ProjectDetectIn, session: Session):
    """
    新增训练集合信息,并创建文件夹
    :param detect_in:
    :param session:
    :return:
    """
    detect = ProjectDetect(**detect_in.dict())
    detect.detect_no = random_str(6)
    detect.detect_version = 0
    detect.detect_status = '0'
    url = os.create_folder(detect_url, detect.detect_no, 'images')
    detect.folder_url = url
    detect = pdc.add_detect(detect, session)
    return detect


def del_detect(detect_id: int, session: Session):
    """
    删除推理集合和推理记录
    :param detect_id:
    :param session:
    :return:
    """
    folder_url = []
    detect = pdc.get_detect_by_id(detect_id, session)
    session.delete(detect)
    folder_url.append(detect.folder_url)
    detect_logs = pdc.get_logs(detect_id, session)
    for log in detect_logs:
        folder_url.append(log.detect_folder_url)
    os.create_folder(folder_url)
    session.commit()


def check_image_name(detect_id: int, files: List[UploadFile], session: Session):
    """
    校验上传的文件名称是否重复
    :param detect_id:
    :param files:
    :param session:
    :return:
    """
    for file in files:
        if not pdc.check_img_name(detect_id, file.filename, session):
            return False, file.filename
    return True, None


def upload_detect_imgs(detect: ProjectDetectOut, files: List[UploadFile], session: Session):
    """
    上传推理集合的照片，保存原图，并生成缩略图
    :param detect:
    :param files:
    :param session:
    :return:
    """
    images = []
    for file in files:
        image = ProjectDetectImg()
        image.detect_id = detect.id
        image.file_name = file.filename
        # 保存原图
        path = os.save_images(detect.folder_url, file=file)
        image.image_url = path
        # 生成缩略图
        thumb_image_url = os.file_path(detect.folder_url, 'thumb', ru.random_str(10) + ".jpg")
        os.create_thumbnail(path, thumb_image_url)
        image.thumb_image_url = thumb_image_url
        images.append(image)
    pdc.add_detect_imgs(images, session)


def del_detect_img(detect_img_id: int, session: Session):
    """
    删除训练集合图片
    :param detect_img_id:
    :param session:
    :return:
    """
    detect_img = session.query(ProjectDetectImg).filter_by(id=detect_img_id).first()
    if detect_img is None:
        return 0
    os.delete_file_if_exists(detect_img.image_url, detect_img.thumb_image_url)
    session.delete(detect_img)
    session.commit()
    return 1


def run_detect_yolo(detect_in: ProjectDetectLogIn, detect: ProjectDetect, train: ProjectTrain, session: Session):
    """
    开始推理
    :param detect:
    :param detect_in:
    :param train:
    :param session:
    :return:
    """
    # 推理版本
    version_path = 'v' + str(detect.detect_version + 1)

    # 权重文件
    pt_url = train.best_pt if detect_in.pt_type == 'best' else train.last_pt

    # 推理集合文件路径
    img_url = detect.folder_url

    out_url = os.file_path(detect_url, detect.detect_no, 'detect')

    # 构建推理记录数据
    detect_log = ProjectDetectLog()
    detect_log.detect_name = detect.detect_name
    detect_log.detect_id = detect.id
    detect_log.detect_version = version_path
    detect_log.train_id = train.id
    detect_log.train_version = train.train_version
    detect_log.pt_type = detect_in.pt_type
    detect_log.pt_url = pt_url
    detect_log.folder_url = img_url
    detect_log.detect_folder_url = out_url
    detect_log = pdc.add_detect_log(detect_log, session)
    return detect_log


async def run_detect_img(weights: str, source: str, project: str, name: str, log_id: int, detect_id: int, session: Session):
    """
    执行yolov5的推理
    :param weights: 权重文件
    :param source: 图片所在文件
    :param project: 推理完成的文件位置
    :param name: 版本名称
    :param log_id: 日志id
    :param detect_id: 推理集合id
    :param session:
    :return:
    """
    yolo_path = os.file_path(yolo_url, 'detect.py')
    room = 'detect_' + str(detect_id)
    await room_manager.send_to_room(room, f"stdout: 模型训练开始，请稍等。。。\n")
    commend = ["python", '-u', yolo_path, "--weights", weights, "--source", source, "--name", name, "--project", project, "--save-txt"]
    # 启动子进程
    with subprocess.Popen(
            commend,
            bufsize=1,  # bufsize=0时，为不缓存；bufsize=1时，按行缓存；bufsize为其他正整数时，为按照近似该正整数的字节数缓存
            shell=False,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,  # 这里可以显示yolov5训练过程中出现的进度条等信息
            text=True,  # 缓存内容为文本，避免后续编码显示问题
            encoding='utf-8',
    ) as process:
        while process.poll() is None:
            line = process.stdout.readline()
            process.stdout.flush()  # 刷新缓存，防止缓存过多造成卡死
            if line != '\n':
                await room_manager.send_to_room(room, line + '\n')

        # 等待进程结束并获取返回码
        return_code = process.wait()
        if return_code != 0:
            pdc.update_detect_status(detect_id, -1, session)
        else:
            await room_manager.send_to_room(room, 'success')
            pdc.update_detect_status(detect_id, 2, session)
            detect_imgs = pdc.get_img_list(detect_id, session)
            detect_log_imgs = []
            for detect_img in detect_imgs:
                detect_log_img = ProjectDetectLogImg()
                detect_log_img.log_id = log_id
                image_url = os.file_path(project, name, detect_img.file_name)
                detect_log_img.image_url = image_url
                detect_log_img.file_name = detect_img.file_name
                detect_log_imgs.append(detect_log_img)
            pdc.add_detect_imgs(detect_log_imgs, session)


async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id: int):
    """
    rtsp 视频流推理
    :param detect_id: 训练集的id
    :param weights_pt: 权重文件
    :param rtsp_url: 视频流地址
    :param data: yaml文件
    :return:
    """
    room = 'detect_rtsp_' + str(detect_id)
    # 选择设备（CPU 或 GPU）
    device = select_device('cpu')

    # 加载模型
    model = DetectMultiBackend(weights_pt, device=device, dnn=False, data=data, fp16=False)

    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size((640, 640), s=stride)  # check image size

    dataset = LoadStreams(rtsp_url, img_size=imgsz, stride=stride, auto=pt, vid_stride=1)
    bs = len(dataset)

    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))

    seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))

    time.sleep(3)# 等待3s，等待websocket进入

    for path, im, im0s, vid_cap, s in dataset:
        if room_manager.rooms.get(room):
            with dt[0]:
                im = torch.from_numpy(im).to(model.device)
                im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
                im /= 255  # 0 - 255 to 0.0 - 1.0
                if len(im.shape) == 3:
                    im = im[None]  # expand for batch dim
                if model.xml and im.shape[0] > 1:
                    ims = torch.chunk(im, im.shape[0], 0)

            # Inference
            with dt[1]:
                if model.xml and im.shape[0] > 1:
                    pred = None
                    for image in ims:
                        if pred is None:
                            pred = model(image, augment=False, visualize=False).unsqueeze(0)
                        else:
                            pred = torch.cat((pred, model(image, augment=False, visualize=False).unsqueeze(0)),
                                             dim=0)
                    pred = [pred, None]
                else:
                    pred = model(im, augment=False, visualize=False)
            # NMS
            with dt[2]:
                pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)

            # Process predictions
            for i, det in enumerate(pred):  # per image
                p, im0, frame = path[i], im0s[i].copy(), dataset.count
                annotator = Annotator(im0, line_width=3, example=str(names))
                if len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                    # Write results
                    for *xyxy, conf, cls in reversed(det):
                        c = int(cls)  # integer class
                        label = None if False else (names[c] if False else f"{names[c]} {conf:.2f}")
                        annotator.box_label(xyxy, label, color=colors(c, True))

                # Stream results
                im0 = annotator.result()
                # 将帧编码为 JPEG
                ret, jpeg = cv2.imencode('.jpg', im0)
                if ret:
                    frame_data = jpeg.tobytes()
                    await room_manager.send_stream_to_room(room, frame_data)
        else:
            break




