"""
Ball detector module using Roboflow Hosted Inference and YOLO v8
"""
import os
import numpy as np
from typing import List, Tuple, Optional, Dict
from inference_sdk import InferenceHTTPClient
import cv2


class BallDetector:
    """
    Detects pickleball balls in video frames using Roboflow pre-trained models
    Implements InferenceSlicer technique for better small object detection
    """

    def __init__(
        self,
        model_id: str = "pickleball-detection-1oqlw/1",
        confidence_threshold: float = 0.4,
        iou_threshold: float = 0.5,
        slice_enabled: bool = True,
        slice_height: int = 320,
        slice_width: int = 320,
        overlap_ratio: float = 0.2
    ):
        """
        Initialize the ball detector

        Args:
            model_id: Roboflow model ID (format: workspace/project/version)
            confidence_threshold: Minimum confidence for detections
            iou_threshold: IoU threshold for NMS
            slice_enabled: Enable frame slicing for better small object detection
            slice_height: Height of each slice
            slice_width: Width of each slice
            overlap_ratio: Overlap ratio between slices
        """
        self.model_id = model_id
        self.confidence_threshold = confidence_threshold
        self.iou_threshold = iou_threshold
        self.slice_enabled = slice_enabled
        self.slice_height = slice_height
        self.slice_width = slice_width
        self.overlap_ratio = overlap_ratio

        # Initialize Roboflow Hosted Inference client
        api_key = os.getenv("ROBOFLOW_API_KEY")

        if not api_key:
            print("✗ ROBOFLOW_API_KEY not set, using YOLO v8 fallback")
            from ultralytics import YOLO
            self.model = YOLO('yolov8n.pt')
            self.use_fallback = True
        else:
            try:
                self.client = InferenceHTTPClient(
                    api_url="https://serverless.roboflow.com",
                    api_key=api_key
                )
                print(f"✓ Initialized Roboflow Hosted Inference: {model_id}")
                self.use_fallback = False
            except Exception as e:
                print(f"✗ Failed to initialize Roboflow client: {e}")
                print("Falling back to YOLO v8 base model for sports ball detection")
                from ultralytics import YOLO
                self.model = YOLO('yolov8n.pt')
                self.use_fallback = True

    def detect(self, frame: np.ndarray) -> List[Dict]:
        """
        Detect balls in a single frame

        Args:
            frame: Input frame (numpy array in BGR format)

        Returns:
            List of detections with format:
            [
                {
                    'bbox': [x1, y1, x2, y2],
                    'confidence': float,
                    'center': [cx, cy]
                },
                ...
            ]
        """
        if self.use_fallback:
            return self._detect_with_yolo(frame)

        if self.slice_enabled:
            return self._detect_with_slicing(frame)
        else:
            return self._detect_single(frame)

    def _detect_single(self, frame: np.ndarray) -> List[Dict]:
        """Detect on full frame without slicing"""
        try:
            # Hosted API expects file path or base64, so we save temp image
            import tempfile
            with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
                cv2.imwrite(tmp.name, frame)
                tmp_path = tmp.name

            try:
                results = self.client.infer(
                    tmp_path,
                    model_id=self.model_id
                )
                return self._parse_results(results)
            finally:
                import os
                os.unlink(tmp_path)

        except Exception as e:
            print(f"Detection error: {e}")
            return []

    def _detect_with_slicing(self, frame: np.ndarray) -> List[Dict]:
        """
        Detect using InferenceSlicer technique
        Divides frame into overlapping tiles for better small object detection
        """
        height, width = frame.shape[:2]
        detections = []

        # Calculate number of slices
        stride_h = int(self.slice_height * (1 - self.overlap_ratio))
        stride_w = int(self.slice_width * (1 - self.overlap_ratio))

        for y in range(0, height, stride_h):
            for x in range(0, width, stride_w):
                # Extract slice
                y_end = min(y + self.slice_height, height)
                x_end = min(x + self.slice_width, width)
                slice_img = frame[y:y_end, x:x_end]

                # Detect on slice
                try:
                    # Hosted API expects file path, save temp slice
                    import tempfile
                    with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
                        cv2.imwrite(tmp.name, slice_img)
                        tmp_path = tmp.name

                    try:
                        results = self.client.infer(
                            tmp_path,
                            model_id=self.model_id
                        )

                        # Parse and adjust coordinates
                        slice_detections = self._parse_results(results)
                        for det in slice_detections:
                            # Adjust bbox coordinates to full frame
                            det['bbox'][0] += x  # x1
                            det['bbox'][1] += y  # y1
                            det['bbox'][2] += x  # x2
                            det['bbox'][3] += y  # y2
                            det['center'][0] += x  # cx
                            det['center'][1] += y  # cy
                            detections.append(det)
                    finally:
                        import os
                        os.unlink(tmp_path)

                except Exception as e:
                    continue

        # Apply NMS to remove duplicate detections from overlapping slices
        detections = self._apply_nms(detections)

        return detections

    def _detect_with_yolo(self, frame: np.ndarray) -> List[Dict]:
        """Fallback detection using ultralytics YOLO"""
        results = self.model(frame, conf=self.confidence_threshold, verbose=False)

        detections = []
        for result in results:
            boxes = result.boxes
            for box in boxes:
                # Filter for sports ball class (class 32 in COCO)
                if int(box.cls[0]) == 32:  # sports ball
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    conf = float(box.conf[0])
                    cx = (x1 + x2) / 2
                    cy = (y1 + y2) / 2

                    detections.append({
                        'bbox': [float(x1), float(y1), float(x2), float(y2)],
                        'confidence': conf,
                        'center': [float(cx), float(cy)]
                    })

        return detections

    def _parse_results(self, results) -> List[Dict]:
        """Parse Roboflow Hosted API results"""
        detections = []

        if not results:
            return detections

        # Hosted API returns dict with 'predictions' key
        predictions = results.get('predictions', []) if isinstance(results, dict) else []

        for pred in predictions:
            # Hosted API format: dict with 'x', 'y', 'width', 'height', 'class', 'confidence'
            if not isinstance(pred, dict):
                continue

            # Extract bbox
            if 'x' in pred and 'y' in pred and 'width' in pred and 'height' in pred:
                cx = pred['x']
                cy = pred['y']
                w = pred['width']
                h = pred['height']
                x1 = cx - w / 2
                y1 = cy - h / 2
                x2 = cx + w / 2
                y2 = cy + h / 2
            else:
                continue

            # Filter for ball class
            class_name = pred.get('class', '')
            if 'ball' not in class_name.lower():
                continue

            conf = pred.get('confidence', 0.0)

            detections.append({
                'bbox': [x1, y1, x2, y2],
                'confidence': conf,
                'center': [cx, cy]
            })

        return detections

    def _apply_nms(self, detections: List[Dict]) -> List[Dict]:
        """Apply Non-Maximum Suppression to remove duplicate detections"""
        if len(detections) == 0:
            return []

        # Extract boxes and scores
        boxes = np.array([det['bbox'] for det in detections])
        scores = np.array([det['confidence'] for det in detections])

        # Apply NMS using OpenCV
        indices = cv2.dnn.NMSBoxes(
            boxes.tolist(),
            scores.tolist(),
            self.confidence_threshold,
            self.iou_threshold
        )

        if len(indices) == 0:
            return []

        # Return filtered detections
        indices = indices.flatten()
        return [detections[i] for i in indices]
