# pylint: disable=too-many-positional-arguments,broad-exception-caught


"""HuggingFace Whisper transcription processor using openai/whisper-large-v3."""

import logging
import os
import tempfile
from typing import Optional

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

from mini_videorag.core.processor import ProcessorResult
from mini_videorag.core.exceptions import (
    VideoLoadError,
    InsufficientDataError,
    ProcessingError,
)
from mini_videorag.processors.transcription.base import BaseTranscriptionProcessor
from mini_videorag.utils.download_video import download_video_from_url
from mini_videorag.utils.audio_utils import video_to_audio


class HFWhisperProcessor(BaseTranscriptionProcessor):
    """
    Transcribes video audio using HuggingFace Whisper models.

    Uses openai/whisper-large-v3 by default, a free open-source model
    that runs locally without API costs. Automatically uses GPU if available.
    """

    def __init__(
        self,
        model_id: str = "openai/whisper-large-v3",
        device: Optional[str] = None,
        torch_dtype: Optional[torch.dtype] = None,
        language: Optional[str] = None,
        task: str = "transcribe",
    ):
        """
        Initialize HuggingFace Whisper processor.

        Args:
            model_id: HuggingFace model ID (default: openai/whisper-large-v3)
            device: Device to use ("cuda:0", "cpu", etc). Auto-detected if None.
            torch_dtype: Torch dtype for model. Auto-detected if None (fp16 for GPU, fp32 for CPU).
            language: Language code for transcription (e.g., "en", "fr"). Auto-detect if None.
            task: Task type - "transcribe" (original language) or "translate" (to English)
        """
        self.model_id = model_id
        self.language = language
        self.task = task
        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

        if torch_dtype is None:
            if torch.cuda.is_available() and "cuda" in self.device:
                self.torch_dtype = torch.float16
            else:
                self.torch_dtype = torch.float32
        else:
            self.torch_dtype = torch_dtype

        self._pipe = None

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

    def _init_pipeline(self):
        """Lazy initialization of Whisper pipeline."""
        if self._pipe is None:
            self.logger.info(f"Loading Whisper model: {self.model_id} on {self.device}")

            # Load model
            model = AutoModelForSpeechSeq2Seq.from_pretrained(
                self.model_id,
                torch_dtype=self.torch_dtype,
                low_cpu_mem_usage=True,
                use_safetensors=True,
            )
            model.to(self.device)

            processor = AutoProcessor.from_pretrained(self.model_id)

            self._pipe = pipeline(
                "automatic-speech-recognition",
                model=model,
                tokenizer=processor.tokenizer,
                feature_extractor=processor.feature_extractor,
                torch_dtype=self.torch_dtype,
                device=self.device,
            )

            self.logger.info("Whisper pipeline initialized successfully")

    def _transcribe_audio(self, audio_path: str) -> str:
        """
        Transcribe audio file using HuggingFace Whisper.

        Args:
            audio_path: Path to audio file

        Returns:
            Transcribed text
        """
        self._init_pipeline()

        generate_kwargs = {}
        if self.language:
            generate_kwargs["language"] = self.language
        if self.task:
            generate_kwargs["task"] = self.task

        result = self._pipe(audio_path, generate_kwargs=generate_kwargs)
        return result["text"].strip()

    def process(self, video_path: str, video_id: str, **kwargs) -> ProcessorResult:
        """
        Transcribe video audio.

        Args:
            video_path: Path to video file
            video_id: Unique identifier
            **kwargs: Optional 'video_url' to download from URL

        Returns:
            ProcessorResult with transcription data
        """
        try:
            self.logger.info(f"Starting HF Whisper transcription for video: {video_id}")

            with tempfile.TemporaryDirectory() as temp_dir:
                if "video_url" in kwargs and not os.path.exists(video_path):
                    video_url = kwargs["video_url"]
                    self.logger.debug(f"Downloading video from: {video_url}")
                    try:
                        download_video_from_url(video_url, video_path)
                    except Exception as e:
                        raise VideoLoadError(f"Failed to download video from URL: {e}") from e

                self.validate_input(video_path)

                audio_path = os.path.join(temp_dir, "audio.mp3")
                self.logger.debug("Extracting audio from video")
                try:
                    video_to_audio(video_path, audio_path)
                except Exception as e:
                    raise InsufficientDataError(
                        f"Failed to extract audio (video may have no audio track): {e}"
                    ) from e

                self.logger.debug("Transcribing audio with HuggingFace Whisper")
                try:
                    transcript = self._transcribe_audio(audio_path)
                except Exception as e:
                    raise ProcessingError(f"Failed to transcribe audio with HF Whisper: {e}") from e

                if self.task == "translate":
                    original_transcript = None
                    english_transcript = transcript
                else:
                    original_transcript = transcript
                    english_transcript = transcript

                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=True,
                    data={
                        "original_transcript": original_transcript,
                        "english_transcript": english_transcript,
                        "language": self.language or "auto-detected",
                        "task": self.task,
                    },
                    metadata={
                        "model_id": self.model_id,
                        "device": self.device,
                        "torch_dtype": str(self.torch_dtype),
                        "audio_extracted": True,
                    },
                )

        except VideoLoadError as e:
            self.logger.error(f"Transcription 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 InsufficientDataError as e:
            self.logger.error(f"Transcription failed - insufficient data: {e}", exc_info=True)
            return ProcessorResult(
                processor_name=self.name,
                video_id=video_id,
                success=False,
                error=f"InsufficientDataError: {e}",
            )
        except ProcessingError as e:
            self.logger.error(f"Transcription failed - processing error: {e}", exc_info=True)
            return ProcessorResult(
                processor_name=self.name,
                video_id=video_id,
                success=False,
                error=f"ProcessingError: {e}",
            )
        except Exception as e:
            self.logger.error(f"Transcription failed - unexpected error: {e}", exc_info=True)
            return ProcessorResult(
                processor_name=self.name,
                video_id=video_id,
                success=False,
                error=f"UnexpectedError: {e}",
            )
