video_detect/router/model_router.py

import os
from pathlib import Path
from service.file_service import save_source_file

from fastapi import APIRouter, HTTPException, UploadFile, File, Form, Query
from mysql.connector import Error as MySQLError

from ds.db import db
from encryption.encrypt_decorator import encrypt_response
from schema.model_schema import (
    ModelResponse,
    ModelListResponse
)
from schema.response_schema import APIResponse
from service.model_service import ALLOWED_MODEL_EXT, MAX_MODEL_SIZE, load_yolo_model

router = APIRouter(prefix="/api/models", tags=["模型管理"])


# 上传模型
@router.post("", response_model=APIResponse, summary="上传YOLO模型（.pt格式）")
@encrypt_response()
async def upload_model(
        name: str = Form(..., description="模型名称"),
        description: str = Form(None, description="模型描述"),
        file: UploadFile = File(..., description=f"YOLO模型文件（.pt、最大{MAX_MODEL_SIZE // 1024 // 1024}MB）")
):
    conn = None
    cursor = None
    try:
        # 校验文件格式
        file_ext = file.filename.split(".")[-1].lower() if "." in file.filename else ""
        if file_ext not in ALLOWED_MODEL_EXT:
            raise HTTPException(
                status_code=400,
                detail=f"仅支持{ALLOWED_MODEL_EXT}格式、当前：{file_ext}"
            )

        # 校验文件大小
        if file.size > MAX_MODEL_SIZE:
            raise HTTPException(
                status_code=400,
                detail=f"文件过大！最大{MAX_MODEL_SIZE // 1024 // 1024}MB、当前{file.size // 1024 // 1024}MB"
            )
        # 保存文件
        file_path = save_source_file(file, "model")

        # 数据库操作
        conn = db.get_connection()
        cursor = conn.cursor(dictionary=True)

        insert_sql = """
            INSERT INTO model (name, path, is_default, description, file_size)
            VALUES (%s, %s, 0, %s, %s)
        """
        cursor.execute(insert_sql, (name, file_path, description, file.size))
        conn.commit()

        # 获取新增记录
        cursor.execute("SELECT * FROM model WHERE id = LAST_INSERT_ID()")
        new_model = cursor.fetchone()
        if not new_model:
            raise HTTPException(status_code=500, detail="上传成功但无法获取记录")

        return APIResponse(
            code=200,
            message=f"模型上传成功",
            data=ModelResponse(**new_model)
        )

    except MySQLError as e:
        if conn:
            conn.rollback()
        raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"服务器错误：{str(e)}") from e
    finally:
        await file.close()
        db.close_connection(conn, cursor)


# 获取模型列表
@router.get("", response_model=APIResponse, summary="获取模型列表（分页）")
@encrypt_response()
async def get_model_list(
        page: int = Query(1, ge=1),
        page_size: int = Query(10, ge=1, le=100),
        name: str = Query(None),
        is_default: bool = Query(None)
):
    conn = None
    cursor = None
    try:
        conn = db.get_connection()
        cursor = conn.cursor(dictionary=True)

        where_clause = []
        params = []
        if name:
            where_clause.append("name LIKE %s")
            params.append(f"%{name}%")
        if is_default is not None:
            where_clause.append("is_default = %s")
            params.append(1 if is_default else 0)

        # 总记录数
        count_sql = "SELECT COUNT(*) AS total FROM model"
        if where_clause:
            count_sql += " WHERE " + " AND ".join(where_clause)
        cursor.execute(count_sql, params)
        total = cursor.fetchone()["total"]

        # 分页数据
        offset = (page - 1) * page_size
        list_sql = "SELECT * FROM model"
        if where_clause:
            list_sql += " WHERE " + " AND ".join(where_clause)
        list_sql += " ORDER BY updated_at DESC LIMIT %s OFFSET %s"
        params.extend([page_size, offset])

        cursor.execute(list_sql, params)
        model_list = cursor.fetchall()

        return APIResponse(
            code=200,
            message=f"获取成功！",
            data=ModelListResponse(
                total=total,
                models=[ModelResponse(**model) for model in model_list]
            )
        )

    except MySQLError as e:
        raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e
    finally:
        db.close_connection(conn, cursor)


# 更换默认模型
@router.put("/{model_id}/set-default", response_model=APIResponse, summary="更换默认模型")
@encrypt_response()
async def set_default_model(
        model_id: int
):
    conn = None
    cursor = None
    try:
        conn = db.get_connection()
        cursor = conn.cursor(dictionary=True)
        conn.autocommit = False

        # 校验目标模型是否存在
        cursor.execute("SELECT * FROM model WHERE id = %s", (model_id,))
        target_model = cursor.fetchone()
        if not target_model:
            raise HTTPException(status_code=404, detail=f"目标模型不存在！")

        # 检查是否已为默认模型
        if target_model["is_default"]:
            return APIResponse(
                code=200,
                message=f"已是默认模型、无需更换",
                data=ModelResponse(**target_model)
            )

        # 数据库事务：更新默认模型状态
        try:
            cursor.execute("UPDATE model SET is_default = 0, updated_at = CURRENT_TIMESTAMP")
            cursor.execute(
                "UPDATE model SET is_default = 1, updated_at = CURRENT_TIMESTAMP WHERE id = %s",
                (model_id,)
            )
            conn.commit()
        except MySQLError as e:
            conn.rollback()
            raise HTTPException(
                status_code=500,
                detail=f"更新默认模型状态失败（已回滚）：{str(e)}"
            ) from e

        # 更新模型
        load_yolo_model()
        # 返回成功响应
        return APIResponse(
            code=200,
            message=f"更换成功",
            data=None
        )

    except MySQLError as e:
        if conn:
            conn.rollback()
        raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e
    finally:
        if conn:
            conn.autocommit = True
        db.close_connection(conn, cursor)


# 路由文件（如 model_router.py）中的删除接口
@router.delete("/{model_id}", response_model=APIResponse, summary="删除模型")
@encrypt_response()
async def delete_model(model_id: int):
    # 1. 正确导入 model_service 中的全局变量（关键修复：变量名匹配）
    from service.model_service import (
        current_yolo_model,
        current_model_absolute_path,
        load_yolo_model  # 用于删除后重新加载模型（可选）
    )

    conn = None
    cursor = None
    try:
        conn = db.get_connection()
        cursor = conn.cursor(dictionary=True)

        # 2. 查询待删除模型信息
        cursor.execute("SELECT * FROM model WHERE id = %s", (model_id,))
        exist_model = cursor.fetchone()
        if not exist_model:
            raise HTTPException(status_code=404, detail=f"模型不存在！")

        # 3. 关键判断：①默认模型不可删 ②正在使用的模型不可删
        if exist_model["is_default"]:
            raise HTTPException(status_code=400, detail="默认模型不可删除！")

        # 计算待删除模型的绝对路径（与 model_service 逻辑一致）
        from service.file_service import get_absolute_path
        del_model_abs_path = get_absolute_path(exist_model["path"])

        # 判断是否正在使用（对比 current_model_absolute_path）
        if current_model_absolute_path and del_model_abs_path == current_model_absolute_path:
            raise HTTPException(status_code=400, detail="该模型正在使用中，禁止删除！")

        # 4. 先删除数据库记录（避免文件删除失败导致数据不一致）
        cursor.execute("DELETE FROM model WHERE id = %s", (model_id,))
        conn.commit()

        # 5. 再删除本地文件（捕获文件删除异常，不影响数据库删除结果）
        extra_msg = ""
        try:
            if os.path.exists(del_model_abs_path):
                os.remove(del_model_abs_path)  # 或用 Path(del_model_abs_path).unlink()
                extra_msg = "（本地文件已同步删除）"
            else:
                extra_msg = "（本地文件不存在，无需删除）"
        except Exception as e:
            extra_msg = f"（本地文件删除失败：{str(e)}）"

        # 6. 若删除后当前模型为空（极端情况），重新加载默认模型（可选优化）
        if current_yolo_model is None:
            try:
                load_yolo_model()
                print(f"[模型删除后] 重新加载默认模型成功")
            except Exception as e:
                print(f"[模型删除后] 重新加载默认模型失败：{str(e)}")

        return APIResponse(
            code=200,
            message=f"模型删除成功！",
            data=None
        )

    except MySQLError as e:
        if conn:
            conn.rollback()
        raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e
    finally:
        db.close_connection(conn, cursor)