"""Brand detection using Hugging Face Falconsai/brand_identification model."""

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.brand_detection.base import BaseBrandDetectionProcessor
from mini_videorag.core.exceptions import (
    ProcessingError,
    ImageProcessingError,
    ImageOpenError,
)
from transformers import pipeline


class HFBrandDetectionProcessor(BaseBrandDetectionProcessor):
    """
    Detect brands and logos using Hugging Face brand identification model.

    Uses Falconsai/brand_identification for logo detection.
    Uses pre-extracted frames from VideoRAG (video_to_images).
    """

    def __init__(
        self,
        brand_model: str = "Falconsai/brand_identification",
        confidence_threshold: float = 0.3,
        device: Optional[str] = None,
    ):
        """
        Initialize brand detection processor.

        Args:
            brand_model: HF model for brand detection (default: Falconsai/brand_identification)
            confidence_threshold: Minimum score for brand detection (default: 0.3)
            device: Device to use ("cuda:0", "mps", "cpu"). Auto-detected if None.
        """
        self.brand_model = brand_model
        self.confidence_threshold = confidence_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._brand_classifier = None

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

    def _init_brand_classifier(self):
        """Lazy initialization of brand classifier."""
        if self._brand_classifier is None:
            self.logger.info(f"Loading brand model: {self.brand_model} on device {self.device}")
            self._brand_classifier = pipeline(
                "image-classification", model=self.brand_model, device=self.device
            )

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

        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 brand 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 brand detection for video: {video_id}")

            self._init_brand_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",
                )

            all_brands = []
            brand_counts = {}
            segments = []

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

                if result.get("brands"):
                    all_brands.extend(result["brands"])

                    for brand in result["brands"]:
                        brand_counts[brand] = brand_counts.get(brand, 0) + 1

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

            unique_brands = list(set(all_brands))

            return ProcessorResult(
                processor_name=self.name,
                video_id=video_id,
                success=True,
                data={
                    "brands_detected": unique_brands,
                    "brand_count": len(all_brands),
                    "brand_counts": brand_counts,
                    "frames_analyzed": len(frame_paths),
                    "frames_with_brands": len(segments),
                },
                segments=segments,
                metadata={
                    "brand_model": self.brand_model,
                    "confidence_threshold": self.confidence_threshold,
                },
            )

        except ProcessingError 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 _detect_brands_in_frame(self, frame_path: str, frame_idx: int) -> Dict:
        """
        Detect brands in a single frame.

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

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

        try:
            predictions = self._brand_classifier(img)

            for pred in predictions:
                label = pred["label"]
                score = pred["score"]

                result["scores"][label] = score

                if score > self.confidence_threshold:
                    result["brands"].append(label)

            if result["brands"]:
                self.logger.debug(f"Frame {frame_idx} brands: {result['brands']}")

        except ImageProcessingError as e:
            raise ProcessingError(
                f"HFBrandDetectionProcessor: Brand detection failed for frame {frame_idx}: {e}"
            ) from e

        return result
