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

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 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_commend(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)


def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str):
    """
    rtsp 视频流推理
    :param weights_pt: 权重文件
    :param rtsp_url: 视频流地址
    :param data: yaml文件
    :return:
    """
    # 选择设备（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))

    for path, im, im0s, vid_cap, s in dataset:
        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
            seen += 1
            p, im0, frame = path[i], im0s[i].copy(), dataset.count
            s += f"{i}: "

            p = Path(p)  # to Path
            s += "{:g}x{:g} ".format(*im.shape[2:])  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            imc = im0.copy() if False else im0  # for save_crop
            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()

                # Print results
                for c in det[:, 5].unique():
                    n = (det[:, 5] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    c = int(cls)  # integer class
                    label = names[c] if False else f"{names[c]}"
                    confidence = float(conf)
                    confidence_str = f"{confidence:.2f}"

                    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()
            if platform.system() == "Linux" and p not in windows:
                windows.append(p)
                cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
                cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
            cv2.imshow(str(p), im0)
            cv2.waitKey(1)  # 1 millisecond