import cv2
import numpy as np
from datetime import datetime
import os

class VisualizationUtils:
    @staticmethod
    def draw_detection_results(
        image: np.ndarray,
        persons: list,
        distances: list = None,
        save_path: str = "output",
        show_confidence: bool = True
    ) -> np.ndarray:
        """
        在图像上绘制检测结果和距离信息
        
        Args:
            image: 原始图像
            persons: 检测到的人物列表，每个元素包含bbox和confidence
            distances: 对应的距离列表（可选）
            save_path: 保存路径
            show_confidence: 是否显示置信度
            
        Returns:
            绘制了检测结果的图像
        """
        # 创建图像副本
        vis_image = image.copy()
        
        # 定义颜色和字体
        BOX_COLOR = (0, 255, 0)  # 绿色边框
        TEXT_COLOR = (255, 255, 255)  # 白色文字
        FONT = cv2.FONT_HERSHEY_SIMPLEX
        FONT_SCALE = 0.6
        THICKNESS = 2
        
        # 遍历每个检测结果
        for idx, person in enumerate(persons):
            bbox = person['bbox']
            conf = person['confidence']
            
            # 绘制边界框
            cv2.rectangle(
                vis_image,
                (bbox[0], bbox[1]),
                (bbox[2], bbox[3]),
                BOX_COLOR,
                THICKNESS
            )
            
            # 准备显示文本
            text_items = []
            if show_confidence:
                text_items.append(f"Conf: {conf:.2f}")
            
            if distances and idx < len(distances):
                dist_meters = distances[idx] / 1000  # 转换为米
                text_items.append(f"Dist: {dist_meters:.2f}m")
            
            # 绘制文本
            text = " | ".join(text_items)
            if text:
                # 获取文本大小
                (text_width, text_height), _ = cv2.getTextSize(
                    text, FONT, FONT_SCALE, THICKNESS
                )
                
                # 绘制文本背景
                cv2.rectangle(
                    vis_image,
                    (bbox[0], bbox[1] - text_height - 10),
                    (bbox[0] + text_width + 10, bbox[1]),
                    BOX_COLOR,
                    -1  # 填充矩形
                )
                
                # 绘制文本
                cv2.putText(
                    vis_image,
                    text,
                    (bbox[0] + 5, bbox[1] - 5),
                    FONT,
                    FONT_SCALE,
                    TEXT_COLOR,
                    THICKNESS
                )
        
        return vis_image
    
    @staticmethod
    def save_detection_image(
        image: np.ndarray,
        save_path: str = "output"
    ) -> str:
        """
        保存检测结果图像
        
        Args:
            image: 要保存的图像
            save_path: 保存路径
            
        Returns:
            保存的文件路径
        """
        # 确保输出目录存在
        os.makedirs(save_path, exist_ok=True)
        
        # 生成文件名（使用时间戳）
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        filename = f"detection_{timestamp}.jpg"
        filepath = os.path.join(save_path, filename)
        
        # 保存图像
        cv2.imwrite(filepath, image)
        
        return filepath