#!/usr/bin/python
# -*- coding: utf-8 -*-
# @version        : 1.0
# @Create Time    : 2025/04/03 10:30
# @File           : views.py
# @IDE            : PyCharm
# @desc           : 路由，视图文件

from utils import os_utils as osu
from core.dependencies import IdList
from utils.websocket_server import room_manager
from . import schemas, crud, params, service, models
from apps.business.train.crud import ProjectTrainDal
from apps.vadmin.auth.utils.current import AllUserAuth
from apps.vadmin.auth.utils.validation.auth import Auth
from utils.response import SuccessResponse, ErrorResponse
from apps.business.deepsort import service as deep_sort_service
from apps.business.project.crud import ProjectInfoDal, ProjectLabelDal

import os
import shutil
import threading
from fastapi.responses import FileResponse
from fastapi import Depends, APIRouter, Form, UploadFile, BackgroundTasks


app = APIRouter()


###########################################################
#    项目推理集合信息
###########################################################
@app.get("/list/{proj_id}", summary="获取推理集合列表")
async def detect_list(
        proj_id: int,
        auth: Auth = Depends(AllUserAuth())):
    datas = await crud.ProjectDetectDal(auth.db).get_datas(
        limit=0,
        v_where=[models.ProjectDetect.project_id == proj_id],
        v_order="asc",
        v_order_field="id",
        v_return_count=False)
    return SuccessResponse(datas)


@app.post("/", summary="创建推理集合")
async def add_detect(
        data: schemas.ProjectDetectIn,
        auth: Auth = Depends(AllUserAuth())):
    detect_dal = crud.ProjectDetectDal(auth.db)
    if await detect_dal.check_name(data.detect_name, data.project_id):
        return ErrorResponse(msg="该项目中存在相同名称的集合")
    await detect_dal.add_detect(data=data, user_id=auth.user.id)
    return SuccessResponse(msg="保存成功")


@app.delete("/", summary="删除项目推理集合信息")
async def delete_detect(
        ids: IdList = Depends(),
        auth: Auth = Depends(AllUserAuth())):
    await crud.ProjectDetectDal(auth.db).delete_detects(ids=ids.ids)
    return SuccessResponse("删除成功")


###########################################################
#    项目推理集合文件信息
###########################################################
@app.get("/files", summary="获取推理集合文件列表")
async def file_list(
        p: params.ProjectDetectFileParams = Depends(),
        auth: Auth = Depends(AllUserAuth())):
    if p.limit > 0:
        datas, count = await crud.ProjectDetectFileDal(auth.db).get_datas(**p.dict(), v_return_count=True)
        return SuccessResponse(datas, count=count)
    else:
        datas = await crud.ProjectDetectFileDal(auth.db).get_datas(**p.dict(), v_return_count=False)
        return SuccessResponse(datas)


@app.post("/files", summary="上传项目推理集合文件")
async def upload_file(
        detect_id: int = Form(...),
        files: list[UploadFile] = Form(...),
        auth: Auth = Depends(AllUserAuth())):
    detect_dal = crud.ProjectDetectDal(auth.db)
    file_dal = crud.ProjectDetectFileDal(auth.db)
    detect_info = await detect_dal.get_data(data_id=detect_id)
    if detect_info is None:
        return ErrorResponse("训练集合查询失败，请刷新后再试")
    if detect_info.file_type != 'rtsp':
        if not any(osu.is_extensions(detect_info.file_type, file.filename) for file in files):
            return ErrorResponse("上传的文件中存在不符合的文件类型")
    await file_dal.add_files(detect_info, files, auth.user.id)
    return SuccessResponse(msg="上传成功")


@app.delete("/files", summary="删除推理集合文件")
async def delete_file(
        ids: IdList = Depends(),
        auth: Auth = Depends(AllUserAuth())):
    await crud.ProjectDetectFileDal(auth.db).delete_files(ids=ids.ids)
    return SuccessResponse("删除成功")


@app.post("/start", summary="开始推理")
async def run_detect_yolo(
        detect_log_in: schemas.ProjectDetectLogIn,
        auth: Auth = Depends(AllUserAuth())):
    detect_dal = crud.ProjectDetectDal(auth.db)
    train_dal = ProjectTrainDal(auth.db)
    detect = await detect_dal.get_data(detect_log_in.detect_id)
    if detect is None:
        return ErrorResponse(msg="训练集合不存在")
    train = await train_dal.get_data(detect_log_in.train_id)
    if train is None:
        return ErrorResponse("训练权重不存在")
    file_count = await crud.ProjectDetectFileDal(auth.db).file_count(detect_log_in.detect_id)
    if file_count == 0 and detect.rtsp_url is None:
        return ErrorResponse("推理集合中没有文件，请先到推理集合中上传推理内容")
    # 判断一下是单纯的推理项目还是跟踪项目
    project_info = await ProjectInfoDal(auth.db).get_data(data_id=detect.project_id)
    if project_info.type_code == 'yolo':
        if detect.file_type == 'img' or detect.file_type == 'video':
            detect_log = await service.before_detect(detect_log_in, detect, train, auth.db, auth.user.id)
            thread_train = threading.Thread(
                target=service.run_detect_folder,
                args=(
                    detect_log.pt_url,
                    detect_log.folder_url,
                    detect_log.detect_folder_url,
                    detect_log.detect_version
                )
            )
            thread_train.start()
            await service.update_sql(
                auth.db,
                detect_log.detect_id,
                detect_log.id,
                detect_log.detect_folder_url,
                detect_log.detect_version
            )
            return SuccessResponse(msg="执行成功", data=file_count)
        elif detect.file_type == 'rtsp':
            room = 'detect_rtsp_' + str(detect.id)
            if not room_manager.rooms.get(room):
                if detect_log_in.pt_type == 'best':
                    weights_pt = train.best_pt
                else:
                    weights_pt = train.last_pt
                thread_train = threading.Thread(
                    target=service.run_detect_rtsp,
                    args=(
                        weights_pt,
                        detect.rtsp_url,
                        room
                    )
                )
                thread_train.start()
            return SuccessResponse(msg="执行成功")
    elif project_info.type_code == 'deepsort':
        room = 'deep_sort_' + str(detect.id)
        if not room_manager.rooms.get(room):
            # 查询项目所属标签，返回两个 id，name一一对应的数组
            label_id_list, label_name_list = await ProjectLabelDal(auth.db).get_label_for_train(project_info.id)
            idx_to_class = {str(i): name for i, name in enumerate(label_name_list)}
            threading_main = threading.Thread(
                target=deep_sort_service.run_deepsort,
                args=(
                    detect.id,
                    train.best_pt,
                    idx_to_class,
                    detect.rtsp_url
                )
            )
            threading_main.start()
        return SuccessResponse(msg="执行成功")


###########################################################
#    项目推理记录信息
###########################################################
@app.get("/logs", summary="获取推理记录列表")
async def logs(
        p: params.ProjectDetectLogParams = Depends(),
        auth: Auth = Depends(AllUserAuth())):
    datas, count = await crud.ProjectDetectLogDal(auth.db).get_datas(**p.dict(), v_return_count=True)
    return SuccessResponse(datas, count=count)


@app.get("/logs/download/{log_id}", summary="下载推理结果")
async def logs_download(
        log_id: int,
        background_tasks: BackgroundTasks,
        auth: Auth = Depends(AllUserAuth())
):
    # 获取日志记录
    detect_log = await crud.ProjectDetectLogDal(auth.db).get_data(data_id=log_id)
    if detect_log is None:
        return ErrorResponse(msg="推理结果查询错误，请刷新页面再试")

    # 检查源文件夹是否存在
    folder_path = os.path.join(detect_log.detect_folder_url, detect_log.detect_version)
    if not os.path.exists(folder_path):
        return ErrorResponse(msg="推理结果尚未生成，请稍后再下载")

    zip_filename = f"{detect_log.detect_name}_{detect_log.detect_version}.zip"
    zip_file_path = osu.zip_folder(folder_path, zip_filename)

    # 后台任务删掉这个
    background_tasks.add_task(cleanup_temp_dir, zip_file_path)

    # 直接返回文件响应
    return FileResponse(path=zip_file_path, filename=os.path.basename(zip_file_path))


def cleanup_temp_dir(temp_dir: str):
    """清理临时目录的后台任务"""
    try:
        shutil.rmtree(temp_dir, ignore_errors=True)
    except Exception as e:
        print(f"清理临时目录失败: {e}")


@app.get("/logs/files", summary="获取项目推理记录文件列表")
async def logs_files(
        p: params.ProjectDetectLogFileParams = Depends(),
        auth: Auth = Depends(AllUserAuth())):
    datas = await crud.ProjectDetectLogFileDal(auth.db).get_datas(**p.dict(), v_return_count=False)
    return SuccessResponse(datas)

