# pylint: disable=redefined-outer-name,reimported,import-outside-toplevel,broad-exception-caught,too-many-positional-arguments,too-many-locals
"""YOLO-based person detection processor."""

import logging
import os
import tempfile
from typing import List, Optional

import cv2
import torch
from moviepy import VideoFileClip
from ultralytics import YOLO

from mini_videorag.core.processor import ProcessorResult, TimeSegment
from mini_videorag.core.exceptions import (
    VideoLoadError,
    ModelInitError,
)
from mini_videorag.processors.person_detection.base import (
    BasePersonDetectionProcessor,
)
from mini_videorag.utils.clip_video import get_frames_from_dir


class YOLOPersonDetectionProcessor(BasePersonDetectionProcessor):
    """
    Detect persons in video frames using YOLO.

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

    def __init__(self, model_path: str = "yolo11n.pt", device: Optional[str] = None):
        """
        Initialize YOLO person detection processor.

        Args:
            model_path: Path to YOLO model file
            device: Device to use ("cuda:0", "mps", "cpu"). Auto-detected if None.
        """
        self.model_path = model_path
        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:
            try:
                self.logger.info(f"Loading YOLO model: {self.model_path} on {self.device}")
                import os

                os.environ["YOLO_VERBOSE"] = "False"
                self.model = YOLO(self.model_path, verbose=False)
                self.model.to(self.device)
            except Exception as e:
                raise ModelInitError(f"Failed to load YOLO model: {e}") from e

    def process(
        self,
        video_path: str,
        video_id: str,
        num_frames: int = 10,
        frames_dir: Optional[str] = None,
        save_annotated: bool = False,
        **kwargs,
    ) -> ProcessorResult:
        """
        Detect persons 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)
            frames_dir: Optional directory with pre-extracted frames (from frame_extraction processor)
            save_annotated: Whether to save annotated frames
            **kwargs: Additional options

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

            self.logger.info(f"Starting person 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:
                total_detections = 0
                frames_with_person = 0
                segments = []

                for frame_path, timestamp, frame_idx in frames_data:
                    detections = self._detect_person_in_frame(frame_path, frame_idx, save_annotated)

                    if detections["person_count"] > 0:
                        frames_with_person += 1
                        total_detections += detections["person_count"]

                        segments.append(
                            TimeSegment(
                                start=timestamp,
                                end=timestamp,
                                metadata={
                                    "person_count": detections["person_count"],
                                    "frame_index": frame_idx,
                                    "confidences": detections["confidences"],
                                },
                            )
                        )

                detection_percentage = (
                    (frames_with_person / len(frames_data) * 100) if frames_data else 0
                )

                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=True,
                    data={
                        "frames_with_person": frames_with_person,
                        "total_frames": len(frames_data),
                        "detection_percentage": round(detection_percentage, 2),
                        "total_detections": total_detections,
                        "has_person": frames_with_person > 0,
                    },
                    segments=segments,
                    metadata={"model": self.model_path, "num_frames": num_frames},
                )
            finally:
                if cleanup_frames:
                    import shutil

                    shutil.rmtree(temp_dir, ignore_errors=True)

        except VideoLoadError as e:
            self.logger.error(f"Person detection failed - video load error: {e}", exc_info=True)
            return ProcessorResult(
                processor_name=self.name,
                video_id=video_id,
                success=False,
                error=f"VideoLoadError: {e}",
            )
        except ModelInitError as e:
            self.logger.error(f"Person detection failed - model init error: {e}", exc_info=True)
            return ProcessorResult(
                processor_name=self.name,
                video_id=video_id,
                success=False,
                error=f"ModelInitError: {e}",
            )
        except Exception as e:
            error_str = str(e)
            if "Cancelled" in error_str or "CancelledError" in str(type(e)):
                self.logger.warning(f"Person detection cancelled for video: {video_id}")
                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=False,
                    error="Activity cancelled",
                )

            self.logger.error(f"Person detection failed - processing error: {e}", exc_info=True)
            return ProcessorResult(
                processor_name=self.name,
                video_id=video_id,
                success=False,
                error=f"ProcessingError: {e}",
            )

    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
        """
        try:
            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)

                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
        except Exception as e:
            raise VideoLoadError(f"Failed to extract frames from video: {e}") from e

    def _detect_person_in_frame(
        self, frame_path: str, frame_idx: int, save_annotated: bool
    ) -> dict:
        """
        Detect persons in a single frame.

        Returns:
            Dict with person_count and confidences
        """
        image = cv2.imread(frame_path)
        if image is None:
            self.logger.warning(f"Could not read frame: {frame_path}")
            return {"person_count": 0, "confidences": []}

        results = self.model(image, verbose=False)
        person_count = 0
        confidences = []

        for result in results:
            boxes = result.boxes
            if boxes is not None:
                person_boxes = boxes[boxes.cls == 0]
                person_count = len(person_boxes)

                if person_count > 0:
                    confidences = person_boxes.conf.tolist()
                    self.logger.debug(f"Frame {frame_idx}: {person_count} person(s) detected")

                    if save_annotated:
                        annotated = result.plot()
                        out_path = os.path.join(
                            os.path.dirname(frame_path),
                            f"annotated_{os.path.basename(frame_path)}",
                        )
                        cv2.imwrite(out_path, annotated)

        return {"person_count": person_count, "confidences": confidences}
