"""Hugging Face-based content safety detection processor."""

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

from PIL import Image
import torch

from mini_videorag.core.processor import ProcessorResult, TimeSegment
from mini_videorag.processors.content_safety.base import BaseContentSafetyProcessor
from mini_videorag.core.exceptions import (
    ImageProcessingError,
    ImageOpenError,
)
from transformers import pipeline


class HFContentSafetyProcessor(BaseContentSafetyProcessor):
    """
    Detect content safety issues using Hugging Face models.

    Uses pre-extracted frames from VideoRAG (video_to_images).
    """

    def __init__(
        self,
        safety_model: str = "Falconsai/nsfw_image_detection",
        check_safety: bool = True,
        safety_threshold: float = 0.5,
        device: Optional[str] = None,
    ):
        """
        Initialize Hugging Face content safety processor.

        Args:
            safety_model: HF model for safety detection (default: Falconsai/nsfw_image_detection)
            check_safety: Enable safety detection (default: True)
            safety_threshold: Minimum score to flag as unsafe (default: 0.5)
            device: Device to use ("cuda:0", "mps", "cpu"). Auto-detected if None.
        """
        self.safety_model = safety_model
        self.check_safety = check_safety
        self.safety_threshold = safety_threshold
        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 = 0
            else:
                self.device = -1
        else:
            if device == "cpu":
                self.device = -1
            elif device == "mps":
                self.device = "mps"
            elif "cuda" in device:
                self.device = int(device.split(":")[-1]) if ":" in device else 0
            else:
                self.device = device

        self._safety_classifier = None

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

    def _init_safety_classifier(self):
        """Lazy initialization of safety classifier."""
        if self._safety_classifier is None and self.check_safety:
            self.logger.info(f"Loading safety model: {self.safety_model} on device {self.device}")
            self._safety_classifier = pipeline(
                "image-classification", model=self.safety_model, device=self.device
            )

    def process(
        self, video_path: str, video_id: str, frames_dir: Optional[str] = None, **kwargs
    ) -> ProcessorResult:
        """
        Analyze video frames for content safety.

        Args:
            video_path: Path to video file (for validation)
            video_id: Unique identifier
            frames_dir: Directory with pre-extracted frames from VideoRAG
            **kwargs: Additional options

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

            if not frames_dir or not os.path.exists(frames_dir):
                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=False,
                    error=f"frames_dir not provided or doesn't exist: {frames_dir}",
                )

            self.logger.info(f"Starting content safety analysis for video: {video_id}")

            if self.check_safety:
                self._init_safety_classifier()

            frame_paths = self._get_frames_from_dir(frames_dir)

            if not frame_paths:
                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=False,
                    error="No frames found in frames_dir",
                )

            safety_issues = []
            segments = []

            for idx, frame_path in enumerate(frame_paths):
                result = self._analyze_frame(frame_path, idx)

                if result.get("safety_issues"):
                    safety_issues.append(
                        {
                            "frame_idx": idx,
                            "issues": result["safety_issues"],
                            "scores": result.get("scores", {}),
                        }
                    )

                    segments.append(
                        TimeSegment(
                            start=idx,
                            end=idx,
                            metadata={
                                "frame_index": idx,
                                "safety_issues": result["safety_issues"],
                                "scores": result.get("scores", {}),
                            },
                        )
                    )

            has_safety_issues = len(safety_issues) > 0

            return ProcessorResult(
                processor_name=self.name,
                video_id=video_id,
                success=True,
                data={
                    "has_safety_issues": has_safety_issues,
                    "safety_issues": safety_issues,
                    "frames_analyzed": len(frame_paths),
                },
                segments=segments,
                metadata={
                    "safety_model": self.safety_model if self.check_safety else None,
                    "check_safety": self.check_safety,
                    "safety_threshold": self.safety_threshold,
                },
            )

        except ImageProcessingError as e:
            return ProcessorResult(
                processor_name=self.name, video_id=video_id, success=False, error=str(e)
            )

    def _get_frames_from_dir(self, frames_dir: str) -> List[str]:
        """Get frame paths from VideoRAG's frames directory."""
        frames = sorted(
            [
                os.path.join(frames_dir, f)
                for f in os.listdir(frames_dir)
                if f.startswith("frame") and f.lower().endswith(".png")
            ]
        )
        return frames

    def _analyze_frame(self, frame_path: str, frame_idx: int) -> Dict:
        """
        Analyze a single frame for safety issues.

        Returns:
            Dict with 'safety_issues' and 'scores'
        """
        try:
            img = Image.open(frame_path)
        except ImageOpenError as e:
            self.logger.warning(f"Failed to open frame {frame_idx}: {e}")
            return {"safety_issues": [], "scores": {}}

        result = {"safety_issues": [], "scores": {}}

        # Safety detection
        if self.check_safety and self._safety_classifier:
            try:
                predictions = self._safety_classifier(img)

                # Process predictions
                for pred in predictions:
                    label = pred["label"].lower()
                    score = pred["score"]

                    result["scores"][label] = score

                    # Flag unsafe content
                    if "nsfw" in label or "unsafe" in label:
                        if score > self.safety_threshold:
                            result["safety_issues"].append(f"{label}_{score:.2f}")

                self.logger.debug(f"Frame {frame_idx} safety: {result['safety_issues']}")
            except ImageProcessingError as e:
                self.logger.warning(f"Safety detection failed for frame {frame_idx}: {e}")

        return result
