"""
YOLO Detection Module

Runs YOLO model inference for field detection.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np

# Import field mappings from single source of truth
from shared.fields import CLASS_NAMES, CLASS_TO_FIELD


@dataclass
class Detection:
    """Represents a single YOLO detection."""
    class_id: int
    class_name: str
    confidence: float
    bbox: tuple[float, float, float, float]  # (x0, y0, x1, y1) in pixels
    page_no: int = 0

    @property
    def x0(self) -> float:
        return self.bbox[0]

    @property
    def y0(self) -> float:
        return self.bbox[1]

    @property
    def x1(self) -> float:
        return self.bbox[2]

    @property
    def y1(self) -> float:
        return self.bbox[3]

    @property
    def center(self) -> tuple[float, float]:
        return ((self.x0 + self.x1) / 2, (self.y0 + self.y1) / 2)

    @property
    def width(self) -> float:
        return self.x1 - self.x0

    @property
    def height(self) -> float:
        return self.y1 - self.y0

    def get_padded_bbox(
        self,
        padding: float = 0.1,
        image_width: float | None = None,
        image_height: float | None = None
    ) -> tuple[float, float, float, float]:
        """Get bbox with padding for OCR extraction."""
        pad_x = self.width * padding
        pad_y = self.height * padding

        x0 = self.x0 - pad_x
        y0 = self.y0 - pad_y
        x1 = self.x1 + pad_x
        y1 = self.y1 + pad_y

        if image_width:
            x0 = max(0, x0)
            x1 = min(image_width, x1)
        if image_height:
            y0 = max(0, y0)
            y1 = min(image_height, y1)

        return (x0, y0, x1, y1)


# CLASS_NAMES and CLASS_TO_FIELD are now imported from shared.fields
# This ensures consistency with the trained YOLO model


class YOLODetector:
    """YOLO model wrapper for field detection."""

    def __init__(
        self,
        model_path: str | Path,
        confidence_threshold: float = 0.5,
        iou_threshold: float = 0.45,
        device: str = 'auto'
    ):
        """
        Initialize YOLO detector.

        Args:
            model_path: Path to trained YOLO model (.pt file)
            confidence_threshold: Minimum confidence for detections
            iou_threshold: IOU threshold for NMS
            device: Device to run on ('auto', 'cpu', 'cuda', 'mps')
        """
        from ultralytics import YOLO

        self.model = YOLO(model_path)
        self.confidence_threshold = confidence_threshold
        self.iou_threshold = iou_threshold
        self.device = device

    def detect(
        self,
        image: str | Path | np.ndarray,
        page_no: int = 0
    ) -> list[Detection]:
        """
        Run detection on an image.

        Args:
            image: Image path or numpy array
            page_no: Page number for reference

        Returns:
            List of Detection objects
        """
        results = self.model.predict(
            source=image,
            conf=self.confidence_threshold,
            iou=self.iou_threshold,
            device=self.device,
            verbose=False
        )

        detections = []

        for result in results:
            boxes = result.boxes
            if boxes is None:
                continue

            for i in range(len(boxes)):
                class_id = int(boxes.cls[i])
                confidence = float(boxes.conf[i])
                bbox = boxes.xyxy[i].tolist()  # [x0, y0, x1, y1]

                class_name = CLASS_NAMES[class_id] if class_id < len(CLASS_NAMES) else f"class_{class_id}"

                detections.append(Detection(
                    class_id=class_id,
                    class_name=class_name,
                    confidence=confidence,
                    bbox=tuple(bbox),
                    page_no=page_no
                ))

        return detections

    def detect_pdf(
        self,
        pdf_path: str | Path,
        dpi: int = 300
    ) -> dict[int, list[Detection]]:
        """
        Run detection on all pages of a PDF.

        Args:
            pdf_path: Path to PDF file
            dpi: Resolution for rendering

        Returns:
            Dict mapping page number to list of detections
        """
        from shared.pdf.renderer import render_pdf_to_images
        from PIL import Image
        import io

        results = {}

        for page_no, image_bytes in render_pdf_to_images(pdf_path, dpi=dpi):
            # Convert bytes to numpy array
            image = Image.open(io.BytesIO(image_bytes))
            image_array = np.array(image)

            detections = self.detect(image_array, page_no=page_no)
            results[page_no] = detections

        return results

    def get_field_name(self, class_name: str) -> str:
        """Convert class name to field name."""
        return CLASS_TO_FIELD.get(class_name, class_name)
