from ultralytics import YOLO
import cv2
import numpy as np

from core.config import settings
from core.logger import logger


class DefectInspector:
    """
    Handles model inference and defect extraction.
    """

    def __init__(self, model_path: str):
        self.model = YOLO(model_path)
        logger.info(f"Model loaded from: {model_path}")

    def inspect_image(self, image: np.ndarray):
        """
        Run inference and extract defect information.
        """

        try:
            results = self.model(
                image,
                conf=settings.CONF_THRESHOLD,
                imgsz=settings.IMAGE_SIZE,
                verbose=False
            )

            result = results[0]
            defects = []

            if result.masks is None:
                return defects

            masks = result.masks.data.cpu().numpy()
            classes = result.boxes.cls.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()

            height, width = image.shape[:2]
            image_area = height * width

            for i, (mask, cls_id) in enumerate(zip(masks, classes)):

                confidence = float(confidences[i])

                if confidence < settings.SECONDARY_CONF_THRESHOLD:
                    continue

                # -------------------------
                # MASK PROCESSING
                # -------------------------
                mask = cv2.resize(mask, (width, height))
                mask = (mask > 0.3).astype("uint8") * 255

                kernel = np.ones((3, 3), np.uint8)
                mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

                # -------------------------
                # CONTOUR EXTRACTION
                # -------------------------
                contours, _ = cv2.findContours(
                    mask,
                    cv2.RETR_EXTERNAL,
                    cv2.CHAIN_APPROX_SIMPLE
                )

                if not contours:
                    continue

                # ✅ PICK LARGEST CONTOUR (FIX)
                largest_contour = max(contours, key=cv2.contourArea)

                # Optional smoothing
                epsilon = 0.01 * cv2.arcLength(largest_contour, True)
                largest_contour = cv2.approxPolyDP(largest_contour, epsilon, True)

                cnt = largest_contour

                area = cv2.contourArea(cnt)

                if area < settings.MIN_DEFECT_AREA:
                    continue

                x, y, w, h = cv2.boundingRect(cnt)

                length = max(w, h)
                width_def = min(w, h)

                area_ratio = area / image_area

                if area_ratio > settings.MAX_AREA_RATIO:
                    continue

                defects.append({
                    "class_id": int(cls_id),
                    "confidence": confidence,
                    "area_pixels": float(area),
                    "length_pixels": float(length),
                    "width_pixels": float(width_def),
                    "area_ratio": float(area_ratio),
                    "bbox": (x, y, w, h),
                    "contour": cnt
                })

            return defects

        except Exception as e:
            logger.error(f"Inference failed: {e}")
            return []