import os
import logging
from ultralytics import YOLO
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)

class ModelService:
    def __init__(self):
        self.models: Dict[str, YOLO] = {}
        self.model_configs = {
            'fire_detection': {
                'path': os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 
                                    'fire_detection', 'models', 'best.pt'),
                'type': 'yolov10',
                'classes': ['Fire', 'Smoke'],
                'labels': {'Fire': '火焰', 'Smoke': '烟雾'},
                'size': '61MB',
                'description': '基于YOLOv10的火灾烟雾检测模型',
                'name': '火灾检测'
            },
            'helmet_detection': {
                'path': os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 
                                    'yolov', 'yolov8n.pt'),
                'type': 'yolov8',
                'classes': ['person', 'helmet'],
                'labels': {'person': '人员', 'helmet': '安全帽'},
                'size': '6MB',
                'description': '基于YOLOv8的安全帽检测模型',
                'name': '安全帽检测'
            },
            'crowd_detection': {
                'path': os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 
                                    'behavior_detection', 'Crowd-Gathering', 'models', 'yolov8l.pt'),
                'type': 'yolov8',
                'classes': ['person'],
                'labels': {'person': '人员'},
                'size': '100MB',
                'description': '基于YOLOv8的人群聚集检测模型',
                'name': '人群检测'
            },
            'smoking_detection': {
                'path': 'PADDLE_DETECTION',  # 特殊标记，表示使用 PaddleDetection
                'type': 'paddle',
                'classes': ['cigarette'],
                'labels': {'cigarette': '香烟'},
                'size': '27MB',
                'description': '基于PP-YOLOE的抽烟检测模型',
                'name': '抽烟检测',
                'docker_image': 'smoking-detection:test',
                'model_dir': 'output_inference/ppyoloe_crn_s_80e_smoking_visdrone'
            }
        }
    
    def get_available_models(self) -> List[Dict]:
        available_models = []
        for model_id, config in self.model_configs.items():
            # 对于 PaddleDetection 模型，不需要检查本地文件
            if config.get('type') == 'paddle':
                # 检查 Docker 是否可用
                import subprocess
                try:
                    result = subprocess.run(
                        ["docker", "image", "inspect", config['docker_image']],
                        capture_output=True,
                        timeout=5
                    )
                    if result.returncode == 0:
                        available_models.append({
                            'id': model_id,
                            'name': config['name'],
                            'description': config['description'],
                            'classes': config['classes'],
                            'labels': config['labels'],
                            'size': config['size'],
                            'type': config['type']
                        })
                    else:
                        logger.warning(f"PaddleDetection Docker 镜像不可用: {config['docker_image']}")
                except Exception as e:
                    logger.warning(f"检查 PaddleDetection Docker 镜像失败: {e}")
            elif os.path.exists(config['path']):
                available_models.append({
                    'id': model_id,
                    'name': config['name'],
                    'description': config['description'],
                    'classes': config['classes'],
                    'labels': config['labels'],
                    'size': config['size'],
                    'type': config['type']
                })
            else:
                logger.warning(f"模型文件不存在: {config['path']}")
        return available_models
    
    async def load_model(self, model_id: str):
        if model_id not in self.model_configs:
            logger.error(f"未知模型ID: {model_id}")
            return None
        
        # 如果已经加载，直接返回
        if model_id in self.models:
            logger.info(f"模型已加载: {model_id}")
            return self.models[model_id]
        
        config = self.model_configs[model_id]
        
        # 处理 PaddleDetection 模型
        if config.get('type') == 'paddle':
            try:
                from .paddle_detection_service import SmokingDetectionModel
                
                logger.info(f"正在加载 PaddleDetection 抽烟检测模型...")
                model = SmokingDetectionModel()
                self.models[model_id] = model
                logger.info(f"PaddleDetection 模型加载成功: {model_id}")
                return model
            except Exception as e:
                logger.error(f"PaddleDetection 模型加载失败: {e}")
                return None
        
        # 处理 YOLO 模型
        model_path = config['path']
        
        if not os.path.exists(model_path):
            logger.error(f"模型文件不存在: {model_path}")
            return None
        
        try:
            logger.info(f"正在加载模型: {model_id} from {model_path}")
            model = YOLO(model_path)
            self.models[model_id] = model
            logger.info(f"模型加载成功: {model_id}")
            return model
        except Exception as e:
            logger.error(f"模型加载失败: {model_id}, 错误: {e}")
            return None
    
    def get_model(self, model_id: str):
        return self.models.get(model_id)
    
    async def unload_model(self, model_id: str) -> bool:
        if model_id in self.models:
            del self.models[model_id]
            logger.info(f"模型已卸载: {model_id}")
            return True
        return False