import os
import numpy as np
import traceback
from ultralytics import YOLO
from typing import Optional


def load_yolo_model(model_path: str) -> Optional[YOLO]:
    """
    加载YOLO模型（支持v5/v8），并校验模型有效性
    :param model_path: 模型文件的绝对路径
    :return: 加载成功返回YOLO模型实例，失败返回None
    """
    try:
        # 加载前的基础信息检查
        print(f"\n[模型工具] 开始加载模型：{model_path}")
        print(f"[模型工具] 文件是否存在：{os.path.exists(model_path)}")
        if os.path.exists(model_path):
            print(f"[模型工具] 文件大小：{os.path.getsize(model_path) / 1024 / 1024:.2f} MB")

        # 强制重新加载模型，避免缓存问题
        model = YOLO(model_path)

        # 兼容性校验：使用numpy空数组测试模型
        dummy_image = np.zeros((640, 640, 3), dtype=np.uint8)

        try:
            # 优先使用新版本参数
            model.predict(
                source=dummy_image,
                imgsz=640,
                conf=0.25,
                verbose=False,
                stream=False
            )
        except Exception as pred_e:
            print(f"[模型工具] 预测校验兼容处理：{str(pred_e)}")
            # 兼容旧版本YOLO参数
            model.predict(
                img=dummy_image,
                imgsz=640,
                conf=0.25,
                verbose=False
            )

        # 验证模型基本属性
        if not hasattr(model, 'names'):
            print("[模型工具] 警告：模型缺少类别名称属性")
        else:
            print(f"[模型工具] 模型包含类别：{list(model.names.values())[:5]}...")  # 显示前5个类别

        print(f"[模型工具] 模型加载成功！")
        return model

    except Exception as e:
        # 详细错误信息输出
        print(f"\n[模型工具] 加载模型失败！路径：{model_path}")
        print(f"[模型工具] 异常类型：{type(e).__name__}")
        print(f"[模型工具] 异常详情：{str(e)}")
        print(f"[模型工具] 堆栈跟踪：\n{traceback.format_exc()}")
        return None