# AI 추론 — YOLOv8 기반 불량(스크래치/이물/흑점/변형) 검출
import os
import numpy as np

from utils.path_helper import BASE_PATH


class Detector:
    class_names = ["스크래치", "이물", "흑점", "변형"]

    def __init__(self):
        self._model     = None
        self.model_path = None

    def load_model(self, model_path: str) -> bool:
        if model_path and not os.path.isabs(model_path):
            model_path = os.path.join(BASE_PATH, model_path)
        try:
            from ultralytics import YOLO  # 지연 로딩 — 앱 시작 시 torch DLL 오류 방지
            self._model     = YOLO(model_path)
            self.model_path = model_path
            print(f"[AI] 모델 로드 성공: {model_path}")
            return True
        except Exception as e:
            print(f"[AI] 모델 로드 실패: {e}")
            self._model = None
            return False

    def is_loaded(self) -> bool:
        return self._model is not None

    def detect(self, image: np.ndarray) -> list:
        """
        image: numpy array (BGR)
        반환: [{"class_id": int, "class_name": str,
                "confidence": float, "bbox": [x1,y1,x2,y2]}, ...]
        """
        if self._model is None:
            return []
        try:
            results = self._model(image, verbose=False)
            detections = []
            for result in results:
                for box in result.boxes:
                    class_id = int(box.cls[0])
                    detections.append({
                        "class_id":   class_id,
                        "class_name": (
                            self.class_names[class_id]
                            if class_id < len(self.class_names)
                            else str(class_id)
                        ),
                        "confidence": float(box.conf[0]),
                        "bbox":       box.xyxy[0].tolist(),
                    })
            return detections
        except Exception as e:
            print(f"[AI] 추론 오류: {e}")
            return []