import os
from dataclasses import dataclass

import cv2
import torch
from ultralytics import YOLO

from ai.plate.detect_rec_plate import det_rec_plate, draw_result
from ai.plate.plate_recognition.plate_rec import init_model
from config.yolo import logger


@dataclass
class PlateInfo:
    plate_no: str
    plate_color: str
    result_img_path: str


class PlateRecognizer:
    _instance = None
    _ready = False

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super(PlateRecognizer, cls).__new__(cls)
        return cls._instance

    def __init__(
        self,
        detect_model_path,
        rec_model_path,
        output_dir="result",
        device="cuda" if torch.cuda.is_available() else "cpu",
    ):
        if hasattr(self, "_initialized") and self._initialized:
            return
        self.device = torch.device(device)
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)

        try:
            # 模型加载
            self.detect_model = YOLO(detect_model_path)
            self.detect_model.to(self.device)
            self.detect_model.eval()
            self.plate_rec_model = init_model(
                self.device, rec_model_path, is_color=True
            )
            self._initialized = True
            self._ready = True
        except Exception as e:
            self.detect_model = None
            self._ready = False
            logger.error(f"❌车牌 YOLO 模型加载失败: {e}")

    def analyze_image(self, image_path):
        img = cv2.imread(image_path)
        if not self._ready or img is None:
            return []

        # 检测与识别
        result_list = det_rec_plate(
            img, self.detect_model, self.plate_rec_model, self.device
        )

        # 可视化 & 保存
        vis_img, _ = draw_result(img, result_list)
        result_img_path = os.path.join(self.output_dir, os.path.basename(image_path))
        cv2.imwrite(result_img_path, vis_img)

        # 构建返回结果
        results = [
            PlateInfo(
                plate_no=r["plate_no"],
                plate_color=r["plate_color"],
                result_img_path=result_img_path,
            )
            for r in result_list
        ]
        return results


# 初始化一次
recognizer = PlateRecognizer(
    detect_model_path="/app/models/sentinel/yolo26s-plate-detect.pt",
    rec_model_path="/app/models/sentinel/plate_rec_color.pth",
    output_dir="result",
)


# base_dir = Path(r"C:\Users\BBIT\Desktop\yolo26-plate-main")
#
# recognizer = PlateRecognizer(
#     detect_model_path=str(base_dir / "weights" / "yolo26s-plate-detect.pt"),
#     rec_model_path=str(base_dir / "weights" / "plate_rec_color.pth"),
#     output_dir="result",
# )