import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import os


class ChineseTextRenderer:
    def __init__(self, font_path=None):
        self.font_cache = {}
        self.font_path = font_path

    def put_text(self, img, text, position, font_size=20, color=(0, 255, 0)):
        """安全的中文文本绘制函数"""
        # 输入验证
        if img is None:
            raise ValueError("输入图像不能为None")

        if not isinstance(img, np.ndarray):
            raise TypeError(f"输入图像必须是numpy数组，实际是{type(img)}")

        if len(img.shape) != 3 or img.shape[2] != 3:
            # 尝试转换
            if len(img.shape) == 2:
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
            elif len(img.shape) == 3 and img.shape[2] == 4:
                img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
            else:
                raise ValueError(f"不支持的图像格式，shape={img.shape}")

        try:
            # BGR转RGB
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            pil_img = Image.fromarray(img_rgb)
            draw = ImageDraw.Draw(pil_img)

            # 获取字体
            font = self._get_font(font_size)

            # 颜色转换：BGR到RGB
            rgb_color = color[::-1]

            # 绘制文本
            draw.text(position, text, font=font, fill=rgb_color)

            # RGB转BGR
            return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

        except Exception as e:
            print(f"中文渲染失败: {e}，使用英文后备")
            # 后备方案：使用OpenCV绘制英文
            cv2.putText(img, text, position, cv2.FONT_HERSHEY_SIMPLEX,
                        font_size / 30, color, 2)
            return img

    def _get_font(self, size):
        """获取字体对象，带缓存"""
        if size in self.font_cache:
            return self.font_cache[size]

        # 查找字体文件
        font_paths = []
        if self.font_path and os.path.exists(self.font_path):
            font_paths.append(self.font_path)

        # 添加常用字体路径
        font_paths.extend([
            "simhei.ttf",
            "msyh.ttc",
            "C:/Windows/Fonts/simhei.ttf",
            "C:/Windows/Fonts/msyh.ttc",
            "C:/Windows/Fonts/Deng.ttf",
            "C:/Windows/Fonts/simsun.ttc",
        ])

        # 遍历所有可能的字体路径
        font = None
        for path in font_paths:
            if os.path.exists(path):
                try:
                    font = ImageFont.truetype(path, size, encoding="utf-8")
                    print(f"加载字体: {path}")
                    break
                except Exception as e:
                    continue

        # 如果找不到字体，尝试创建默认字体
        if font is None:
            try:
                font = ImageFont.load_default()
                # 调整默认字体大小
                font = ImageFont.truetype("arial.ttf", size)
            except:
                # 最后手段：使用PIL的默认字体
                font = ImageFont.load_default()

        self.font_cache[size] = font
        return font


def ensure_image_valid(img, default_size=(640, 480)):
    """确保图像有效，如果无效则创建默认图像"""
    if img is None or not isinstance(img, np.ndarray) or img.size == 0:
        print("创建默认图像...")
        img = np.zeros((default_size[0], default_size[1], 3), dtype=np.uint8)
        img.fill(50)
        # 添加一些参考线
        h, w = img.shape[:2]
        cv2.line(img, (0, 0), (w, h), (0, 255, 0), 1)
        cv2.line(img, (w, 0), (0, h), (0, 255, 0), 1)
        cv2.circle(img, (w // 2, h // 2), min(w, h) // 4, (255, 0, 0), 2)

    return img


# 使用示例
def run_yolo_with_chinese():
    # 初始化中文渲染器
    renderer = ChineseTextRenderer()

    # 模拟YOLO检测结果
    # 注意：这里假设你已经有检测结果
    detections = [
        {"bbox": [100, 100, 200, 300], "conf": 0.95, "class": "person"},
        {"bbox": [300, 150, 450, 350], "conf": 0.88, "class": "car"},
    ]

    # 中英文类别映射
    class_map = {
        "person": "人",
        "car": "汽车",
        "bicycle": "自行车",
        "dog": "狗",
        "cat": "猫",
        "chair": "椅子",
        "bottle": "瓶子"
    }

    # 读取图像或创建默认图像
    img_path = "test.jpg"
    if os.path.exists(img_path):
        img = cv2.imread(img_path)
    else:
        # 用户选择图像文件
        import tkinter as tk
        from tkinter import filedialog

        root = tk.Tk()
        root.withdraw()
        img_path = filedialog.askopenfilename(
            title="选择图像文件",
            filetypes=[("Image files", "*.jpg *.jpeg *.png *.bmp *.tiff")]
        )

        if img_path:
            img = cv2.imread(img_path)
        else:
            print("未选择文件，创建测试图像")
            img = None

    # 确保图像有效
    img = ensure_image_valid(img)

    # 绘制检测结果
    for det in detections:
        x1, y1, x2, y2 = det["bbox"]
        conf = det["conf"]
        cls_en = det["class"]
        cls_cn = class_map.get(cls_en, cls_en)

        # 绘制边界框
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # 绘制中文标签
        label = f"{cls_cn}: {conf:.2f}"
        img = renderer.put_text(img, label, (x1, max(y1 - 20, 10)),
                                font_size=15, color=(0, 255, 0))

    # 添加标题
    img = renderer.put_text(img, "YOLO检测结果", (10, 30),
                            font_size=25, color=(255, 255, 0))

    # 显示结果
    cv2.imshow("YOLO Detection with Chinese", img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 保存结果
    output_path = "detection_result.jpg"
    cv2.imwrite(output_path, img)
    print(f"结果已保存: {output_path}")


if __name__ == "__main__":
    run_yolo_with_chinese()