# app/cloud/obstacle_detector_client.py (新文件)
import logging
import os
import cv2
import numpy as np
import torch
from threading import Semaphore
from contextlib import contextmanager
from ultralytics import YOLOE
from typing import List, Dict, Any

# Day 20: Numba 多核加速
try:
    from numba_utils import count_mask_pixels, compute_mask_stats, bitwise_and_count, warmup as numba_warmup
    NUMBA_ENABLED = True
except ImportError:
    NUMBA_ENABLED = False

logger = logging.getLogger(__name__)

# --- GPU/CPU & AMP 配置 (从 blindpath 工作流迁移而来，保持一致) ---
DEVICE = os.getenv("AIGLASS_DEVICE", "cuda:0")
if DEVICE.startswith("cuda") and not torch.cuda.is_available():
    logger.warning(f"AIGLASS_DEVICE={DEVICE} 但未检测到 CUDA，将回退到 CPU")
    DEVICE = "cpu"
IS_CUDA = DEVICE.startswith("cuda")

AMP_POLICY = os.getenv("AIGLASS_AMP", "fp16").lower()
if AMP_POLICY not in ("bf16", "fp16", "off"):
    AMP_POLICY = "fp16"
AMP_DTYPE = torch.bfloat16 if AMP_POLICY == "bf16" else (torch.float16 if AMP_POLICY == "fp16" else None)

# --- GPU 并发限流 (从 blindpath 工作流迁移而来，保持一致) ---
# Day 20: 增加默认槽位从 2 到 4，RTX 3090 可以处理更多并发
GPU_SLOTS = int(os.getenv("AIGLASS_GPU_SLOTS", "4"))
_gpu_slots = Semaphore(GPU_SLOTS)

try:
    torch.backends.cudnn.benchmark = True
except Exception:
    pass


@contextmanager
def gpu_infer_slot():
    """统一管理 GPU 并发限流 + inference_mode + AMP autocast"""
    with _gpu_slots:
        if IS_CUDA and AMP_POLICY != "off":
            # 新式接口：torch.amp.autocast(device_type='cuda', dtype=...)
            with torch.inference_mode(), torch.amp.autocast(device_type='cuda', dtype=AMP_DTYPE):
                yield
        else:
            with torch.inference_mode():
                yield


class ObstacleDetectorClient:
    def __init__(self, model_path: str = 'model/yoloe-11l-seg.pt'):
        self.model = None
        self.whitelist_embeddings = None
        self.WHITELIST_CLASSES = [
            'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'animal', 'scooter', 'stroller', 'dog',
            'pole', 'post', 'column', 'pillar', 'stanchion', 'bollard', 'utility pole',
            'telegraph pole', 'light pole', 'street pole', 'signpost', 'support post',
            'vertical post', 'bench', 'chair', 'potted plant', 'hydrant', 'cone', 'stone', 'box'
        ]
        # COCO 类别白名单 - TensorRT 模式下用于后处理过滤
        # 从 COCO 80 类中筛选出可能构成障碍物的类别
        self.COCO_WHITELIST = {
            'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck',  # 交通
            'dog', 'cat', 'horse', 'cow', 'sheep',  # 动物
            'bench', 'chair', 'potted plant', 'fire hydrant', 'stop sign',  # 街道设施
            'parking meter', 'suitcase', 'backpack', 'umbrella', 'handbag',  # 物品
            'sports ball', 'skateboard', 'surfboard', 'bottle', 'cup',  # 可能障碍
        }
        try:
            # Day 20: 优先使用 TensorRT 引擎
            try:
                from model_utils import get_best_model_path, is_tensorrt_engine
                model_path = get_best_model_path(model_path)
            except ImportError:
                def is_tensorrt_engine(p): return p.endswith('.engine')
            
            logger.info(f"正在加载 YOLOE 障碍物模型: {model_path}")
            self.model = YOLOE(model_path)
            
            # Day 20: TensorRT 引擎不需要 .to() 和 .fuse()
            if is_tensorrt_engine(model_path):
                logger.info(f"TensorRT 引擎已加载，跳过 .to() 和 .fuse()")
                # TensorRT 引擎不支持 get_text_pe，跳过白名单特征计算
                self.whitelist_embeddings = None
                logger.info("TensorRT 模式：跳过白名单特征预计算")
            else:
                self.model.to(DEVICE)
                self.model.fuse()
                logger.info(f"YOLOE 障碍物模型加载成功，使用设备: {DEVICE}")

                logger.info("正在为 YOLOE 预计算白名单文本特征...")
                if IS_CUDA and AMP_DTYPE is not None:
                    with torch.inference_mode(), torch.amp.autocast(device_type='cuda', dtype=AMP_DTYPE):
                        self.whitelist_embeddings = self.model.get_text_pe(self.WHITELIST_CLASSES)
                else:
                    self.whitelist_embeddings = self.model.get_text_pe(self.WHITELIST_CLASSES)
                logger.info("YOLOE 特征预计算完成。")
        except Exception as e:
            logger.error(f"YOLOE 模型加载或特征计算失败: {e}", exc_info=True)
            raise
    def tensor_to_numpy_mask(mask_tensor):
        """安全地将各种类型的张量转换为 numpy 掩码"""
        # 处理不同的数据类型
        if mask_tensor.dtype in (torch.bfloat16, torch.float16):
            mask_tensor = mask_tensor.float()
        
        # 转换为 numpy
        mask = mask_tensor.cpu().numpy()
        
        # 确保是二值掩码
        if mask.max() <= 1.0:
            mask = (mask > 0.5).astype(np.uint8) * 255
        else:
            mask = mask.astype(np.uint8)
        
        return mask 
    def detect(self, image: np.ndarray, path_mask: np.ndarray = None) -> List[Dict[str, Any]]:
        """
        利用白名单作为提示词寻找障碍物。
        如果提供了 path_mask，则执行与路径相关的空间过滤。
        如果 path_mask 为 None，则进行全局检测。
        """
        if self.model is None:
            return []

        H, W = image.shape[:2]
        
        # TensorRT 模式下没有 embeddings，跳过 set_classes
        # 此时模型会使用默认的 COCO 类别进行检测
        if self.whitelist_embeddings is not None:
            try:
                self.model.set_classes(self.WHITELIST_CLASSES, self.whitelist_embeddings)
            except Exception as e:
                logger.error(f"设置 YOLOE 提示词失败: {e}")
                return []

        conf_thr = float(os.getenv("AIGLASS_OBS_CONF", "0.25"))
        # Day 22 优化: 动态输入尺寸和FP16加速
        imgsz = int(os.getenv("AIGLASS_OBS_IMGSZ", "480"))  # 从默认640降低
        use_half = os.getenv("AIGLASS_OBS_HALF", "1") == "1"
        
        with gpu_infer_slot():
            results = self.model.predict(
                image, 
                verbose=False, 
                conf=conf_thr,
                imgsz=imgsz,      # 使用较小的输入尺寸
                half=use_half     # FP16 半精度加速
            )

        if not (results and results[0].masks):
            return []

        # --- 过滤与后处理 (逻辑与 blindpath 工作流保持一致) ---
        final_obstacles = []
        num_masks = len(results[0].masks.data)
        num_boxes = len(results[0].boxes.cls) if getattr(results[0].boxes, "cls", None) is not None else 0

        for i, mask_tensor in enumerate(results[0].masks.data):
            if i >= num_boxes: continue

            # 【修复】处理 BFloat16 类型的掩码
            # 先转换为 float32，避免 numpy 不支持 BFloat16 的问题
            if mask_tensor.dtype == torch.bfloat16:
                mask_tensor = mask_tensor.float()
            
            # 转换为 numpy 数组
            mask = mask_tensor.cpu().numpy()
            
            # 处理概率掩码（值在0-1之间）或二值掩码
            if mask.max() <= 1.0:
                # 概率掩码，需要二值化
                mask = (mask > 0.5).astype(np.uint8) * 255
            else:
                # 已经是二值掩码
                mask = mask.astype(np.uint8)
            
            mask = cv2.resize(mask, (W, H), interpolation=cv2.INTER_NEAREST)
            
            # Day 20: 使用 Numba 多核加速计算 mask 统计信息
            if NUMBA_ENABLED:
                stats = compute_mask_stats(mask)
                area = stats['area']
                center_x = stats['center_x']
                center_y = stats['center_y']
                min_y, max_y = stats['bbox'][1], stats['bbox'][3]
            else:
                area = int(np.sum(mask > 0))
                y_coords, x_coords = np.where(mask > 0)
                if len(y_coords) == 0:
                    continue
                center_x = float(np.mean(x_coords))
                center_y = float(np.mean(y_coords))
                min_y, max_y = int(np.min(y_coords)), int(np.max(y_coords))

            # 尺寸过滤：太大的物体（如整片地面）通常是误识别
            if (area / (H * W)) > 0.7: continue
            if area == 0: continue

            # 空间过滤：如果提供了 path_mask，则只保留路径上的障碍物
            if path_mask is not None:
                # Day 20: 使用 Numba 加速交集计算
                if NUMBA_ENABLED:
                    intersection_area = bitwise_and_count(mask, path_mask)
                else:
                    intersection_area = int(np.sum(cv2.bitwise_and(mask, path_mask) > 0))
                # 必须与路径有足够的重叠
                if intersection_area < 100 or (intersection_area / area) < 0.01:
                    continue

            cls_id = int(results[0].boxes.cls[i])
            class_names_map = results[0].names
            class_name = "Unknown"
            if isinstance(class_names_map, dict):
                # 如果是字典，使用 .get() 方法
                class_name = class_names_map.get(cls_id, "Unknown")
            elif isinstance(class_names_map, list) and 0 <= cls_id < len(class_names_map):
                # 如果是列表，通过索引安全地获取
                class_name = class_names_map[cls_id]
            
            # TensorRT 模式下使用 COCO 白名单过滤
            # 只保留可能构成障碍物的类别
            if self.whitelist_embeddings is None:  # TensorRT 模式
                if class_name.lower().strip() not in self.COCO_WHITELIST:
                    continue  # 跳过非白名单类别

            final_obstacles.append({
                'name': class_name.strip(),
                'mask': mask,
                'area': area,
                'area_ratio': area / (H * W),
                'center_x': center_x,
                'center_y': center_y,
                'bottom_y_ratio': max_y / H
            })

        return final_obstacles