"""
Object Detection Service
Uses YOLO for fast and accurate object detection
"""
import logging
from typing import List, Dict, Any, Optional
from io import BytesIO
import numpy as np
from PIL import Image
import cv2

from app.core.config import settings

logger = logging.getLogger(__name__)


class DetectionService:
    """Object detection using YOLO"""

    def __init__(self):
        self.model = None
        self.model_loaded = False

    async def load_models(self):
        """Load detection models"""
        try:
            from ultralytics import YOLO
            model_path = settings.YOLO_MODEL
            self.model = YOLO(model_path)
            self.model_loaded = True
            logger.info(f"YOLO model loaded: {model_path}")
        except Exception as e:
            logger.warning(f"Failed to load YOLO model: {e}")
            self.model_loaded = False

    async def detect_objects(
        self,
        image_bytes: bytes,
        confidence_threshold: float = 0.5,
        classes: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        """
        Detect objects in an image

        Args:
            image_bytes: Image data as bytes
            confidence_threshold: Minimum confidence for detections
            classes: Filter to specific classes

        Returns:
            Detection results
        """
        if not self.model_loaded:
            await self.load_models()

        if not self.model_loaded:
            return {
                "error": "Model not available",
                "detections": []
            }

        # Load image
        image = Image.open(BytesIO(image_bytes))
        image_np = np.array(image)

        # Run detection
        results = self.model(image_np, conf=confidence_threshold)

        detections = []
        for result in results:
            boxes = result.boxes
            for box in boxes:
                class_id = int(box.cls[0])
                class_name = self.model.names[class_id]

                # Filter by classes if specified
                if classes and class_name not in classes:
                    continue

                x1, y1, x2, y2 = box.xyxy[0].tolist()
                confidence = float(box.conf[0])

                detections.append({
                    "class": class_name,
                    "confidence": round(confidence, 3),
                    "bbox": {
                        "x1": int(x1),
                        "y1": int(y1),
                        "x2": int(x2),
                        "y2": int(y2),
                        "width": int(x2 - x1),
                        "height": int(y2 - y1)
                    }
                })

        # Group by class
        class_counts = {}
        for det in detections:
            cls = det["class"]
            class_counts[cls] = class_counts.get(cls, 0) + 1

        return {
            "total_detections": len(detections),
            "detections": detections,
            "class_counts": class_counts,
            "image_size": {"width": image.width, "height": image.height},
            "confidence_threshold": confidence_threshold
        }

    async def detect_and_count(
        self,
        image_bytes: bytes,
        target_class: str,
        confidence_threshold: float = 0.5
    ) -> Dict[str, Any]:
        """
        Count specific objects in an image

        Args:
            image_bytes: Image data
            target_class: Class to count
            confidence_threshold: Minimum confidence

        Returns:
            Count results
        """
        result = await self.detect_objects(
            image_bytes,
            confidence_threshold,
            classes=[target_class]
        )

        return {
            "target_class": target_class,
            "count": result["class_counts"].get(target_class, 0),
            "detections": result["detections"],
            "confidence_threshold": confidence_threshold
        }

    async def detect_with_annotations(
        self,
        image_bytes: bytes,
        confidence_threshold: float = 0.5
    ) -> Dict[str, Any]:
        """
        Detect objects and return annotated image

        Args:
            image_bytes: Image data
            confidence_threshold: Minimum confidence

        Returns:
            Detection results with annotated image bytes
        """
        result = await self.detect_objects(image_bytes, confidence_threshold)

        # Load and annotate image
        image = Image.open(BytesIO(image_bytes))
        image_np = np.array(image)

        # Convert RGB to BGR for OpenCV
        if len(image_np.shape) == 3 and image_np.shape[2] == 3:
            image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
        else:
            image_bgr = image_np

        # Draw bounding boxes
        colors = self._generate_colors(len(result["detections"]))

        for i, det in enumerate(result["detections"]):
            bbox = det["bbox"]
            color = colors[i % len(colors)]

            # Draw rectangle
            cv2.rectangle(
                image_bgr,
                (bbox["x1"], bbox["y1"]),
                (bbox["x2"], bbox["y2"]),
                color,
                2
            )

            # Draw label
            label = f"{det['class']} {det['confidence']:.2f}"
            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)

            cv2.rectangle(
                image_bgr,
                (bbox["x1"], bbox["y1"] - label_size[1] - 10),
                (bbox["x1"] + label_size[0], bbox["y1"]),
                color,
                -1
            )

            cv2.putText(
                image_bgr,
                label,
                (bbox["x1"], bbox["y1"] - 5),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (255, 255, 255),
                2
            )

        # Convert back to RGB and to bytes
        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        annotated_image = Image.fromarray(image_rgb)

        buffer = BytesIO()
        annotated_image.save(buffer, format="JPEG", quality=90)
        annotated_bytes = buffer.getvalue()

        import base64
        result["annotated_image_base64"] = base64.b64encode(annotated_bytes).decode()

        return result

    def _generate_colors(self, n: int) -> List[tuple]:
        """Generate distinct colors for visualization"""
        import colorsys

        colors = []
        for i in range(n):
            hue = i / max(n, 1)
            rgb = colorsys.hsv_to_rgb(hue, 0.8, 0.9)
            colors.append(tuple(int(c * 255) for c in rgb))

        return colors if colors else [(0, 255, 0)]

    def get_available_classes(self) -> List[str]:
        """Get list of classes the model can detect"""
        if self.model_loaded:
            return list(self.model.names.values())
        return []


# Singleton instance
detection_service = DetectionService()