"""
PaddleDetection 抽烟检测服务适配器
使用本地 PaddlePaddle 环境直接调用模型（无需 Docker）
"""

# 禁用 PIR API 以支持旧版模型格式（必须在任何导入之前设置）
import os
os.environ['FLAGS_enable_pir_api'] = '0'

import cv2
import numpy as np
import logging
import threading
import time
import sys
from typing import Dict, List, Optional
from pathlib import Path

logger = logging.getLogger(__name__)


class PaddleDetectionService:
    """PaddleDetection 服务适配器（本地模式）"""
    
    def __init__(self):
        self.model_name = "smoking_detection"
        self.threshold = 0.1
        self._lock = threading.Lock()
        
        # 本地环境配置
        project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
        self.paddle_dir = os.path.join(project_root, "third-party", "paddle-inference")
        self.model_dir = os.path.join(project_root, "models", "smoking_detection_paddle")
        
        # 检测器实例（延迟加载）
        self._detector = None
        self._detector_initialized = False
        
        self.available = True
        logger.info(f"本地 PaddlePaddle 模式已启用")
        logger.info(f"模型目录: {self.model_dir}")
        logger.info(f"使用服务器虚拟环境中的 PaddlePaddle")
        logger.info(f"PaddlePaddle 目录: {self.paddle_dir}")
        
        # 禁用 PIR API 以支持旧版模型格式（必须在初始化前设置）
        os.environ['FLAGS_enable_pir_api'] = '0'
        
        # 检测系统架构
        import platform
        self.platform_info = platform.uname()
        self.is_apple_silicon = self.platform_info.machine in ('arm64', 'aarch64') and self.platform_info.system == 'Darwin'
        
        if self.is_apple_silicon:
            logger.info("✅ 检测到 Apple Silicon (ARM64) 架构")
            logger.info("✅ 使用本地 PaddlePaddle 环境获得最佳性能")
            logger.info("✅ 相比 Docker 方式性能提升 5-10 倍")
        
        try:
            self._initialize_environment()
        except Exception as e:
            logger.error(f"初始化环境失败: {e}")
            self.available = False
    
    def _initialize_environment(self):
        """初始化本地 PaddlePaddle 环境"""
        try:
            # 添加 PaddleDetection 部署路径
            paddle_detection_path = self.paddle_dir
            if paddle_detection_path not in sys.path:
                sys.path.insert(0, paddle_detection_path)
                logger.info(f"✅ 添加 PaddleDetection 路径: {paddle_detection_path}")
            
            # 检查模型目录是否存在
            if not os.path.exists(self.model_dir):
                raise Exception(f"模型目录不存在: {self.model_dir}")
            
            # 检查必要文件
            required_files = ['model.pdmodel', 'model.pdiparams', 'infer_cfg.yml']
            for file in required_files:
                file_path = os.path.join(self.model_dir, file)
                if not os.path.exists(file_path):
                    raise Exception(f"模型文件不存在: {file}")
            
            logger.info("✅ 环境检查通过")
            
            # 预加载检测器（可选，用于首次检测预热）
            try:
                self._get_detector()
                logger.info("✅ 检测器预加载成功")
            except Exception as e:
                logger.warning(f"检测器预加载失败，将在首次使用时初始化: {e}")
            
        except Exception as e:
            logger.error(f"环境初始化失败: {e}")
            raise
    
    def _get_detector(self):
        """获取检测器实例（单例模式）"""
        if self._detector is None or not self._detector_initialized:
            try:
                # 设置环境变量以支持旧版模型格式
                os.environ['FLAGS_enable_pir_api'] = '0'
                
                # 添加 PaddleDetection 路径（直接使用 self.paddle_dir）
                if self.paddle_dir not in sys.path:
                    sys.path.insert(0, self.paddle_dir)
                    logger.info(f"添加 PaddleDetection 路径: {self.paddle_dir}")
                
                # 导入 PaddleDetection 模块
                from infer import Detector, PredictConfig
                
                # 创建检测器
                self._detector = Detector(
                    model_dir=self.model_dir,
                    device='CPU',
                    run_mode='paddle',
                    batch_size=1,
                    output_dir='output',
                    threshold=self.threshold
                )
                
                self._detector_initialized = True
                logger.info("✅ PaddlePaddle 检测器初始化成功")
                
            except Exception as e:
                logger.error(f"检测器初始化失败: {e}")
                raise
        
        return self._detector
    
    def detect_image(self, image: np.ndarray, threshold: float = None) -> Dict:
        """
        检测图片中的抽烟行为（本地模式）
        
        Args:
            image: OpenCV 图片 (BGR格式)
            threshold: 置信度阈值，如果为 None 则使用默认值
            
        Returns:
            检测结果字典
        """
        if threshold is None:
            threshold = self.threshold
        
        if not self.available:
            return {
                'success': False,
                'message': 'PaddleDetection 服务不可用',
                'detections': [],
                'stats': None
            }
        
        try:
            with self._lock:
                start_time = time.time()
                
                # 确保检测器已初始化
                detector = self._get_detector()
                
                # 准备输入图片
                if not isinstance(image, np.ndarray):
                    raise Exception(f"不支持的图片类型: {type(image)}")
                
                if len(image.shape) == 2:  # 灰度图转 BGR
                    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
                elif image.shape[2] == 4:  # RGBA 转 BGR
                    image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)
                
                # 执行推理
                inference_start = time.time()
                
                # 使用 PaddleDetection API 进行推理
                results = detector.predict_image(
                    [image], 
                    visual=False, 
                    save_results=False
                )
                
                inference_time = time.time() - inference_start
                logger.info(f"推理耗时: {inference_time:.3f}s")
                
                # 解析检测结果
                detections = self._parse_detection_results(results, threshold)
                
                total_time = time.time() - start_time
                logger.info(f"检测总耗时: {total_time:.3f}s")
                
                return {
                    'success': True,
                    'message': '检测完成',
                    'detections': detections,
                    'stats': {
                        'total_detections': len(detections),
                        'model_used': 'ppyoloe_crn_s_80e_smoking_visdrone',
                        'threshold': threshold,
                        'processing_time': round(total_time, 3),
                        'inference_time': round(inference_time, 3)
                    }
                }
        
        except Exception as e:
            import traceback
            logger.error(f"检测失败: {e}")
            logger.error(f"错误堆栈: {traceback.format_exc()}")
            
            # 重置检测器状态以允许重试
            self._detector_initialized = False
            
            return {
                'success': False,
                'message': f'检测失败: {e}',
                'detections': [],
                'stats': None
            }
    
    def _parse_detection_results(self, results: Dict, threshold: float) -> List[Dict]:
        """解析 PaddleDetection 返回的检测结果"""
        detections = []
        
        try:
            if results and 'boxes' in results:
                boxes = results['boxes']
                
                if boxes is not None and len(boxes) > 0:
                    for box in boxes:
                        # 解析检测结果格式: [class_id, score, x1, y1, x2, y2]
                        if len(box) >= 6:
                            class_id = int(box[0])
                            confidence = float(box[1])
                            x1, y1, x2, y2 = float(box[2]), float(box[3]), float(box[4]), float(box[5])
                            
                            # 过滤低置信度检测
                            if confidence >= threshold:
                                detections.append({
                                    'class': 'cigarette',
                                    'label': '香烟',
                                    'confidence': round(confidence, 3),
                                    'bbox': [int(x1), int(y1), int(x2), int(y2)]
                                })
        
        except Exception as e:
            logger.error(f"解析检测结果失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
        
        return detections
    
    def get_performance_info(self) -> Dict:
        """获取性能信息"""
        return {
            'mode': 'local',
            'environment': 'PaddlePaddle',
            'model_dir': self.model_dir,
            'apple_silicon': self.is_apple_silicon,
            'detector_loaded': self._detector_initialized,
            'available': self.available
        }


# 兼容性包装，保持与 YOLO 模型相同的接口
class SmokingDetectionModel:
    """抽烟检测模型包装器，兼容 YOLO 接口"""
    
    def __init__(self):
        self.service = PaddleDetectionService()
        self.names = {0: 'cigarette'}
    
    def __call__(self, image, conf=0.1, iou=0.45, verbose=False):
        """
        模拟 YOLO 模型的调用接口
        
        Args:
            image: OpenCV 图片
            conf: 置信度阈值
            iou: IoU 阈值
            verbose: 是否输出详细信息
            
        Returns:
            模拟 YOLO 结果的对象
        """
        result = self.service.detect_image(image, threshold=conf)
        
        return [PaddleDetectionResult(result, self.names)]


class PaddleDetectionResult:
    """模拟 YOLO 检测结果对象"""
    
    def __init__(self, detection_result: Dict, names: Dict):
        self.detection_result = detection_result
        self.names = names
        
        self.boxes = self._create_boxes()
    
    def _create_boxes(self):
        """创建模拟的 boxes 对象"""
        detections = self.detection_result.get('detections', [])
        
        if not detections:
            return MockBoxes([])
        
        xyxy = []
        conf = []
        cls = []
        
        for det in detections:
            xyxy.append(det['bbox'])
            conf.append(det['confidence'])
            cls.append(0)
        
        return MockBoxes(xyxy, conf, cls)


class MockBoxes:
    """模拟 YOLO boxes 对象"""
    
    def __init__(self, xyxy_list, conf_list=None, cls_list=None):
        try:
            import torch
            use_torch = True
        except ImportError:
            use_torch = False
        
        if xyxy_list and len(xyxy_list) > 0:
            if use_torch:
                self.xyxy = torch.tensor(xyxy_list, dtype=torch.float32)
                self.conf = torch.tensor(conf_list, dtype=torch.float32).reshape(-1, 1)
                self.cls = torch.tensor(cls_list, dtype=torch.int64).reshape(-1, 1)
            else:
                self.xyxy = np.array(xyxy_list, dtype=np.float32)
                self.conf = np.array(conf_list, dtype=np.float32).reshape(-1, 1)
                self.cls = np.array(cls_list, dtype=np.int64).reshape(-1, 1)
        else:
            if use_torch:
                self.xyxy = torch.empty((0, 4), dtype=torch.float32)
                self.conf = torch.empty((0, 1), dtype=torch.float32)
                self.cls = torch.empty((0, 1), dtype=torch.int64)
            else:
                self.xyxy = np.array([]).reshape(0, 4)
                self.conf = np.array([]).reshape(0, 1)
                self.cls = np.array([]).reshape(0, 1)
        
        self._use_torch = use_torch
    
    def __iter__(self):
        for i in range(len(self.xyxy)):
            yield MockBox(
                self.xyxy[i],
                self.conf[i][0] if len(self.conf) > i else 0.0,
                self.cls[i][0] if len(self.cls) > i else 0
            )
    
    def __len__(self):
        return len(self.xyxy)
    
    def cpu(self):
        return self
    
    def numpy(self):
        if self._use_torch:
            if len(self.xyxy) > 0:
                return (
                    self.xyxy.numpy(),
                    self.conf.numpy(),
                    self.cls.numpy()
                )
            else:
                return (
                    np.array([]).reshape(0, 4),
                    np.array([]).reshape(0, 1),
                    np.array([], dtype=np.int64).reshape(0, 1)
                )
        else:
            return (
                self.xyxy,
                self.conf,
                self.cls
            )


class MockBox:
    """模拟单个 YOLO box 对象"""
    
    def __init__(self, xyxy, conf, cls):
        try:
            import torch
            use_torch = True
        except ImportError:
            use_torch = False
        
        if use_torch:
            if isinstance(xyxy, torch.Tensor):
                self.xyxy = xyxy
            else:
                self.xyxy = torch.tensor(xyxy, dtype=torch.float32)
        else:
            if isinstance(xyxy, np.ndarray):
                self.xyxy = xyxy
            else:
                self.xyxy = np.array(xyxy, dtype=np.float32)
        
        self.conf = conf
        self.cls = cls