# pylint: disable=import-outside-toplevel, too-many-positional-arguments,too-many-locals
"""YOLO-based object detection processor."""
import logging
import os
import tempfile
from typing import Dict, List, Optional
from collections import Counter
import cv2
import torch
from moviepy import VideoFileClip
from ultralytics import YOLO
from mini_videorag.core.exceptions import ImageProcessingError, ProcessingError
from mini_videorag.core.processor import ProcessorResult, TimeSegment
from mini_videorag.processors.object_detection.base import (
    BaseObjectDetectionProcessor,
)
from mini_videorag.utils.clip_video import get_frames_from_dir


class YOLOObjectDetectionProcessor(BaseObjectDetectionProcessor):
    """
    Detect objects in video frames using YOLO.

    Uses Ultralytics YOLO to detect objects in video frames
    and returns detection statistics with timestamps.
    """

    def __init__(
        self,
        model_path: str = "models/yolo11n.pt",
        num_frames: int = 10,
        confidence_threshold: float = 0.25,
        device: Optional[str] = None,
    ):
        """
        Initialize YOLO object detection processor.

        Args:
            model_path: Path to YOLO model file
            num_frames: Number of frames to sample (default: 10)
            confidence_threshold: Minimum confidence for detections (default: 0.25)
            device: Device to use ("cuda:0", "mps", "cpu"). Auto-detected if None.
        """
        self.model_path = model_path
        self.num_frames = num_frames
        self.confidence_threshold = confidence_threshold
        self.model = None
        self.logger = logging.getLogger(__name__)

        if device is None:
            if torch.backends.mps.is_available():
                self.device = "mps"
            elif torch.cuda.is_available():
                self.device = "cuda:0"
            else:
                self.device = "cpu"
        else:
            self.device = device

    @property
    def provider(self) -> str:
        """Provider name for this implementation."""
        return "yolo"

    def _init_model(self):
        """Lazy initialization of YOLO model."""
        if self.model is None:
            self.logger.info(f"Loading YOLO model: {self.model_path} on {self.device}")
            os.environ["YOLO_VERBOSE"] = "False"
            self.model = YOLO(self.model_path, verbose=False)
            self.model.to(self.device)

    def process(
        self,
        video_path: str,
        video_id: str,
        num_frames: Optional[int] = None,
        frames_dir: Optional[str] = None,
        confidence_threshold: Optional[float] = None,
        **kwargs,
    ) -> ProcessorResult:
        """
        Detect objects in video frames.

        Args:
            video_path: Path to video file
            video_id: Unique identifier
            num_frames: Number of frames to sample (only used if frames_dir not provided, uses instance default if not provided)
            frames_dir: Optional directory with pre-extracted frames (from frame_extraction processor)
            confidence_threshold: Minimum confidence for detections (uses instance default if not provided)
            **kwargs: Additional options

        Returns:
            ProcessorResult with detection data
        """
        try:
            self.validate_input(video_path)
            self._init_model()

            num_frames = num_frames if num_frames is not None else self.num_frames
            if confidence_threshold is None:
                confidence_threshold = self.confidence_threshold

            self.logger.info(f"Starting object detection for video: {video_id}")

            if frames_dir and os.path.exists(frames_dir):
                self.logger.info(f"Using pre-extracted frames from: {frames_dir}")
                frames_data = get_frames_from_dir(frames_dir, include_frame_index=True)
                cleanup_frames = False
            else:
                self.logger.info(f"Extracting {num_frames} frames on-the-fly")
                temp_dir = tempfile.mkdtemp()
                frames_data = self._extract_frames(video_path, temp_dir, num_frames)
                cleanup_frames = True

            try:
                all_objects = []
                segments = []

                for frame_path, timestamp, frame_idx in frames_data:
                    detections = self._detect_objects_in_frame(
                        frame_path, frame_idx, confidence_threshold
                    )

                    if detections["objects"]:
                        all_objects.extend(detections["objects"])
                        segments.append(
                            TimeSegment(
                                start=timestamp,
                                end=timestamp,
                                metadata={
                                    "frame_index": frame_idx,
                                    "objects": detections["objects"],
                                    "object_count": len(detections["objects"]),
                                },
                            )
                        )

                object_counter = Counter(all_objects)
                unique_objects = list(object_counter.keys())

                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=True,
                    data={
                        "unique_objects": unique_objects,
                        "object_counts": dict(object_counter),
                        "total_detections": len(all_objects),
                        "frames_processed": len(frames_data),
                    },
                    segments=segments,
                    metadata={
                        "model": self.model_path,
                        "num_frames": num_frames,
                        "confidence_threshold": confidence_threshold,
                    },
                )
            finally:
                if cleanup_frames:
                    import shutil

                    shutil.rmtree(temp_dir, ignore_errors=True)

        except ImageProcessingError as e:
            raise ProcessingError(
                f"YOLOObjectDetectionProcessor: Object detection failed: {e}"
            ) from e
        except Exception as e:
            error_str = str(e)
            if "Cancelled" in error_str or "CancelledError" in str(type(e)):
                self.logger.warning(f"Object detection cancelled for video: {video_id}")
                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=False,
                    error="Activity cancelled",
                )
            raise

    def _extract_frames(self, video_path: str, output_dir: str, num_frames: int) -> List[tuple]:
        """
        Extract evenly spaced frames from video.

        Returns:
            List of (frame_path, timestamp, index) tuples
        """
        frames_data = []
        video = VideoFileClip(video_path)
        duration = video.duration

        for i in range(num_frames):
            timestamp = (i / (num_frames - 1)) * duration if num_frames > 1 else 0
            frame = video.get_frame(timestamp)

            # Save frame
            frame_path = os.path.join(output_dir, f"frame_{i:04d}.jpg")
            cv2.imwrite(frame_path, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))

            frames_data.append((frame_path, timestamp, i))

        video.close()
        return frames_data

    def _detect_objects_in_frame(
        self, frame_path: str, frame_idx: int, confidence_threshold: float
    ) -> Dict:
        """Detect objects in a single frame.

        Returns:
            Dict with list of detected object class names
        """
        image = cv2.imread(frame_path)
        if image is None:
            self.logger.warning(f"Could not read frame: {frame_path}")
            return {"objects": []}

        results = self.model(image, conf=confidence_threshold, verbose=False)
        detected_objects = []

        for result in results:
            boxes = result.boxes
            if boxes is not None and len(boxes) > 0:
                for cls_id in boxes.cls:
                    class_name = self.model.names[int(cls_id)]
                    detected_objects.append(class_name)

                if detected_objects:
                    self.logger.debug(
                        f"Frame {frame_idx}: {len(detected_objects)} object(s) detected: {set(detected_objects)}"
                    )

        return {"objects": detected_objects}
