# pylint: disable=broad-exception-caught, too-many-locals,too-many-positional-arguments, import-outside-toplevel

"""EasyOCR text extraction processor with language detection and stopword filtering."""

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

import cv2
import easyocr
import torch
from moviepy import VideoFileClip

from mini_videorag.core.processor import ProcessorResult, TimeSegment
from mini_videorag.core.exceptions import (
    VideoLoadError,
    ModelInitError,
)
from mini_videorag.processors.ocr.base import BaseOCRProcessor
from mini_videorag.utils.clip_video import get_frames_from_dir
from mini_videorag.processors.ocr.text_utils import (
    clean_and_tokenize,
    fix_suspicious_tokens,
    filter_ocr_tokens,
    detect_language,
)

logger = logging.getLogger(__name__)


class EasyOCRProcessor(BaseOCRProcessor):
    """
    Extract text from video frames using EasyOCR.

    Enhanced with:
    - Language detection
    - Stopword filtering
    - Token normalization
    - Suspicious token fixing (word splitting)
    """

    def __init__(
        self,
        languages: List[str] = None,
        enable_stopword_filtering: bool = True,
        enable_token_fixing: bool = True,
        confidence_threshold: float = 0.5,
        device: Optional[str] = None,
    ):
        """
        Initialize EasyOCR processor.

        Args:
            languages: List of language codes (e.g., ['en', 'fr'])
            enable_stopword_filtering: Enable stopword removal
            enable_token_fixing: Enable wordninja token splitting
            confidence_threshold: Minimum confidence for OCR results
            device: Device to use ("cuda:0", "mps", "cpu"). Auto-detected if None.
        """
        self.languages = languages or ["en"]
        self.enable_stopword_filtering = enable_stopword_filtering
        self.enable_token_fixing = enable_token_fixing
        self.confidence_threshold = confidence_threshold
        self.reader = 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 "easyocr"

    def _init_reader(self):
        """Lazy initialization of EasyOCR reader."""
        if self.reader is None:
            try:
                self.logger.info(
                    f"Initializing EasyOCR with languages: {self.languages} on {self.device}"
                )
                gpu = self.device != "cpu"
                self.reader = easyocr.Reader(self.languages, gpu=gpu)
            except Exception as e:
                raise ModelInitError(f"Failed to initialize EasyOCR reader: {e}") from e

    def process(
        self,
        video_path: str,
        video_id: str,
        num_frames: int = 10,
        frames_dir: Optional[str] = None,
        **kwargs,
    ) -> ProcessorResult:
        """
        Extract text from video frames with enhanced processing.

        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)
            **kwargs: Additional options

        Returns:
            ProcessorResult with OCR text data including language detection
        """
        try:
            self.validate_input(video_path)
            self._init_reader()

            self.logger.info(f"Starting OCR 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=False)
                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_text = []
                all_raw_text = []
                filtered_tokens = []
                segments = []

                for frame_path, timestamp in frames_data:
                    text_results = self.reader.readtext(frame_path)

                    frame_texts = []
                    raw_frame_text = []

                    for detection in text_results:
                        bbox, text, confidence = detection
                        if confidence > self.confidence_threshold:
                            frame_texts.append(
                                {"text": text, "confidence": confidence, "bbox": bbox}
                            )
                            all_text.append(text)
                            raw_frame_text.append(text)

                    if raw_frame_text:
                        frame_combined = " ".join(raw_frame_text)
                        all_raw_text.append(frame_combined)

                        tokens = clean_and_tokenize(frame_combined)

                        if self.enable_token_fixing:
                            tokens = fix_suspicious_tokens(tokens)

                        if self.enable_stopword_filtering:
                            tokens = filter_ocr_tokens(tokens, self.languages[0])

                        filtered_tokens.extend(tokens)

                    if frame_texts:
                        segments.append(
                            TimeSegment(
                                start=timestamp,
                                end=timestamp,
                                metadata={
                                    "texts": frame_texts,
                                    "frame_path": frame_path,
                                },
                            )
                        )

                combined_text = " ".join(all_raw_text)
                detected_language = detect_language(combined_text) if combined_text else None

                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=True,
                    data={
                        "all_text": all_text,
                        "filtered_tokens": filtered_tokens,
                        "detected_language": detected_language,
                        "text_count": len(all_text),
                        "filtered_token_count": len(filtered_tokens),
                        "frames_processed": len(frames_data),
                        "frames_with_text": len(segments),
                    },
                    segments=segments,
                    metadata={
                        "languages": self.languages,
                        "num_frames": num_frames,
                        "detected_language": detected_language,
                        "stopword_filtering_enabled": self.enable_stopword_filtering,
                        "token_fixing_enabled": self.enable_token_fixing,
                        "confidence_threshold": self.confidence_threshold,
                    },
                )
            finally:
                if cleanup_frames:
                    import shutil

                    shutil.rmtree(temp_dir, ignore_errors=True)

        except VideoLoadError as e:
            self.logger.error(f"OCR 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"OCR 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"OCR cancelled for video: {video_id}")
                return ProcessorResult(
                    processor_name=self.name,
                    video_id=video_id,
                    success=False,
                    error="Activity cancelled",
                )

            self.logger.error(f"OCR 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) 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))

            video.close()
            return frames_data
        except Exception as e:
            raise VideoLoadError(f"Failed to extract frames from video: {e}") from e
