#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
视频检测系统配置文件
"""

import os
from pathlib import Path
from ultralytics import YOLO

class DetectionConfig:
    """检测系统配置类"""
    
    # 模型配置
    MODEL_PATH = "models/best.pt"
    CONFIDENCE_THRESHOLD = 0.5
    
    # 动态模型信息
    _model_classes = None
    _model_colors = None
    
    # 文件夹配置
    INPUT_FOLDER = "input_videos"
    OUTPUT_FOLDER = "output_frames"
    LOG_FOLDER = "logs"
    
    # 视频处理配置
    SUPPORTED_VIDEO_FORMATS = {'.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv', '.webm'}
    PROGRESS_INTERVAL = 100  # 每处理多少帧显示一次进度
    MIN_FRAME_INTERVAL = 1.0  # 检测帧的最小时间间隔（秒）
    
    # 输出配置
    SAVE_ORIGINAL_FRAMES = False  # 是否保存原始帧
    SAVE_ANNOTATED_FRAMES = True  # 是否保存标注帧
    SAVE_DETECTION_JSON = True  # 是否保存检测结果JSON
    
    # 图像质量配置
    IMAGE_QUALITY = 95  # JPEG质量 (1-100)
    
    # 默认道路损伤类别映射 (中文标签) - 作为备用
    DEFAULT_CLASS_NAMES_CN = {
        0: "纵向裂缝",
        1: "横向裂缝", 
        2: "网状裂缝",
        3: "坑洞",
        4: "白线模糊"
    }
    
    # 默认类别颜色配置 (BGR格式) - 所有类别使用相同的绿色
    DEFAULT_CLASS_COLORS = {
        0: (0, 255, 0),    # 绿色
        1: (0, 255, 0),    # 绿色
        2: (0, 255, 0),    # 绿色
        3: (0, 255, 0),    # 绿色
        4: (0, 255, 0)     # 绿色
    }
    
    @classmethod
    def get_model_path(cls):
        """获取模型文件的绝对路径"""
        current_dir = Path(__file__).parent
        model_path = current_dir / cls.MODEL_PATH
        return str(model_path.resolve())
    
    @classmethod
    def create_folders(cls):
        """创建必要的文件夹"""
        folders = [cls.INPUT_FOLDER, cls.OUTPUT_FOLDER, cls.LOG_FOLDER]
        for folder in folders:
            Path(folder).mkdir(parents=True, exist_ok=True)
    
    @classmethod
    def load_model_classes(cls, model_path):
        """从模型文件动态加载类别信息"""
        try:
            model = YOLO(model_path)
            cls._model_classes = model.names
            # 为所有类别生成绿色
            cls._model_colors = {i: (0, 255, 0) for i in range(len(model.names))}
            return True
        except Exception as e:
            print(f"加载模型类别信息失败: {e}")
            return False
    
    @classmethod
    def get_class_name_cn(cls, class_id):
        """获取类别的中文名称"""
        # 优先使用动态加载的模型类别
        if cls._model_classes is not None:
            original_name = cls._model_classes.get(class_id)
            if original_name:
                # 如果有对应的中文映射则使用，否则使用原始英文名称
                return cls.DEFAULT_CLASS_NAMES_CN.get(class_id, original_name)
            return f"未知类别_{class_id}"
        # 回退到默认中文类别
        return cls.DEFAULT_CLASS_NAMES_CN.get(class_id, f"未知类别_{class_id}")
    
    @classmethod
    def get_class_color(cls, class_id):
        """获取类别对应的颜色 - 始终返回绿色"""
        if cls._model_colors is not None:
            return cls._model_colors.get(class_id, (0, 255, 0))
        return cls.DEFAULT_CLASS_COLORS.get(class_id, (0, 255, 0))
    
    @classmethod
    def get_all_classes(cls):
        """获取所有类别信息"""
        if cls._model_classes is not None:
            return cls._model_classes
        return cls.DEFAULT_CLASS_NAMES_CN

# 全局配置实例
config = DetectionConfig()