"""
Text and graphics overlay operations for FMUS-VID.

This module provides functionality for adding text and graphics to videos.
"""

from typing import Union, Tuple, Optional, Dict, Any
from pathlib import Path
import os
import logging
from PIL import Image, ImageDraw, ImageFont
import cv2
import numpy as np
import textwrap

logger = logging.getLogger(__name__)

class TextOverlay:
    """Class for handling text overlays in videos."""

    def __init__(self, text: str, position: Tuple[int, int],
                 font: str = "Arial", size: int = 24,
                 color: Union[str, Tuple[int, int, int]] = "white",
                 start_time: float = 0, duration: Optional[float] = None,
                 fade_in: float = 0, fade_out: float = 0):
        """
        Initialize a text overlay.

        Args:
            text: Text to display
            position: (x, y) coordinates for placement
            font: Font name or path
            size: Font size in pixels
            color: Text color (name or RGB tuple)
            start_time: When to start showing the text (seconds)
            duration: How long to show the text (seconds, None = until end)
            fade_in: Fade-in duration (seconds)
            fade_out: Fade-out duration (seconds)
        """
        self.text = text
        self.position = position
        self.font = font
        self.size = size
        self.color = color
        self.start_time = start_time
        self.duration = duration
        self.fade_in = fade_in
        self.fade_out = fade_out

        # Try to load the font
        try:
            self._font_obj = ImageFont.truetype(font, size)
        except IOError:
            logger.warning(f"Font '{font}' not found, using default font")
            self._font_obj = ImageFont.load_default()

    def render(self, frame: Image.Image, time: float) -> Image.Image:
        """
        Render text on a frame.

        Args:
            frame: PIL Image to render on
            time: Current video time in seconds

        Returns:
            Frame with rendered text
        """
        # Check if text should be visible at this time
        if time < self.start_time:
            return frame

        if self.duration is not None and time > self.start_time + self.duration:
            return frame

        # Calculate opacity based on fade effects
        opacity = 1.0

        if self.fade_in > 0 and time < self.start_time + self.fade_in:
            opacity = (time - self.start_time) / self.fade_in

        if self.fade_out > 0 and self.duration is not None:
            fade_start = self.start_time + self.duration - self.fade_out
            if time > fade_start:
                opacity = 1.0 - (time - fade_start) / self.fade_out

        opacity = max(0.0, min(1.0, opacity))

        if opacity == 0:
            return frame

        # Create a copy of the frame
        result = frame.copy()
        draw = ImageDraw.Draw(result)

        # Draw the text with opacity
        if opacity < 1.0:
            # Create a temporary image for the text
            text_overlay = Image.new('RGBA', frame.size, (0, 0, 0, 0))
            text_draw = ImageDraw.Draw(text_overlay)

            # Draw text on the overlay
            text_draw.text(self.position, self.text,
                         font=self._font_obj, fill=(*self._get_color(), int(255 * opacity)))

            # Composite the text overlay onto the frame
            result = Image.alpha_composite(result.convert('RGBA'), text_overlay)
        else:
            # Draw text directly for full opacity
            draw.text(self.position, self.text, font=self._font_obj, fill=self.color)

        return result

    def _get_color(self) -> Tuple[int, int, int]:
        """Convert color to RGB tuple."""
        if isinstance(self.color, tuple):
            return self.color
        return ImageColor.getrgb(self.color)

class GraphicsOverlay:
    """Class for handling graphics overlays in videos."""

    def __init__(self, image_path: Union[str, Path], position: Tuple[int, int],
                 size: Optional[Tuple[int, int]] = None, opacity: float = 1.0,
                 start_time: float = 0, duration: Optional[float] = None,
                 fade_in: float = 0, fade_out: float = 0):
        """
        Initialize a graphics overlay.

        Args:
            image_path: Path to image file
            position: (x, y) coordinates for placement
            size: Optional (width, height) to resize image
            opacity: Base opacity (1.0 = fully opaque)
            start_time: When to start showing the graphic (seconds)
            duration: How long to show the graphic (seconds, None = until end)
            fade_in: Fade-in duration (seconds)
            fade_out: Fade-out duration (seconds)
        """
        self.image_path = Path(image_path)
        self.position = position
        self.size = size
        self.base_opacity = opacity
        self.start_time = start_time
        self.duration = duration
        self.fade_in = fade_in
        self.fade_out = fade_out

        # Load and prepare the image
        self._load_image()

    def _load_image(self):
        """Load and prepare the overlay image."""
        if not self.image_path.is_file():
            raise FileNotFoundError(f"Image not found: {self.image_path}")

        # Load the image
        self.image = Image.open(self.image_path)

        # Convert to RGBA if needed
        if self.image.mode != 'RGBA':
            self.image = self.image.convert('RGBA')

        # Resize if size specified
        if self.size:
            self.image = self.image.resize(self.size, Image.LANCZOS)

    def render(self, frame: Image.Image, time: float) -> Image.Image:
        """
        Render graphic on a frame.

        Args:
            frame: PIL Image to render on
            time: Current video time in seconds

        Returns:
            Frame with rendered graphic
        """
        # Check if graphic should be visible at this time
        if time < self.start_time:
            return frame

        if self.duration is not None and time > self.start_time + self.duration:
            return frame

        # Calculate opacity based on fade effects
        opacity = self.base_opacity

        if self.fade_in > 0 and time < self.start_time + self.fade_in:
            opacity *= (time - self.start_time) / self.fade_in

        if self.fade_out > 0 and self.duration is not None:
            fade_start = self.start_time + self.duration - self.fade_out
            if time > fade_start:
                opacity *= 1.0 - (time - fade_start) / self.fade_out

        opacity = max(0.0, min(1.0, opacity))

        if opacity == 0:
            return frame

        # Create a copy of the frame
        result = frame.copy()

        # Apply opacity to the overlay image
        if opacity < 1.0:
            overlay = self.image.copy()
            r, g, b, a = overlay.split()
            a = a.point(lambda i: int(i * opacity))
            overlay = Image.merge('RGBA', (r, g, b, a))
        else:
            overlay = self.image

        # Paste the overlay onto the frame
        result.paste(overlay, self.position, overlay)

        return result

# Add SubtitleOverlay class for real-time subtitle rendering

class SubtitleOverlay:
    """
    Apply subtitle overlays to video frames.

    This class provides optimized methods for adding subtitles to video frames
    in real-time streaming scenarios.
    """

    def __init__(
        self,
        position: str = "bottom",
        font_path: Optional[str] = None,
        font_size: int = 28,
        font_color: Tuple[int, int, int] = (255, 255, 255),
        background_color: Tuple[int, int, int] = (0, 0, 0),
        background_opacity: float = 0.5,
        padding: int = 10,
        max_width: Optional[int] = None,
        margin: int = 20,
        shadow: bool = True,
        shadow_offset: int = 2,
        shadow_color: Tuple[int, int, int] = (0, 0, 0),
        line_spacing: int = 5,
        text_align: str = "center",
        char_width_factor: float = 0.6  # Approximate character width as a factor of font size
    ):
        """
        Initialize subtitle overlay.

        Args:
            position: Position of subtitle ('top', 'bottom', 'center')
            font_path: Path to font file (None for default)
            font_size: Font size in points
            font_color: RGB color tuple for text
            background_color: RGB color tuple for background
            background_opacity: Opacity of background (0.0 to 1.0)
            padding: Padding around text in pixels
            max_width: Maximum width of subtitle text (None for frame width - 2*margin)
            margin: Margin from edges in pixels
            shadow: Whether to add text shadow
            shadow_offset: Shadow offset in pixels
            shadow_color: RGB color tuple for shadow
            line_spacing: Space between lines in pixels
            text_align: Text alignment ('left', 'center', 'right')
            char_width_factor: Approximate character width as a factor of font size
        """
        self.position = position
        self.font_path = font_path
        self.font_size = font_size
        self.font_color = font_color
        self.background_color = background_color
        self.background_opacity = background_opacity
        self.padding = padding
        self.max_width = max_width
        self.margin = margin
        self.shadow = shadow
        self.shadow_offset = shadow_offset
        self.shadow_color = shadow_color
        self.line_spacing = line_spacing
        self.text_align = text_align
        self.char_width_factor = char_width_factor

        # Internal state
        self._font = None
        self._pil_font = None
        self._text_cache = {}  # Cache rendered text
        self._last_frame_size = None

        # Initialize font
        self._initialize_font()

    def _initialize_font(self):
        """Initialize font for text rendering."""
        try:
            # Try to use FreeType font (OpenCV 3.x+)
            if self.font_path and os.path.exists(self.font_path):
                self._font = cv2.freetype.createFreeType2()
                self._font.loadFontData(self.font_path, 0)
            else:
                # Use a PIL font as backup
                try:
                    import PIL.ImageFont
                    if self.font_path and os.path.exists(self.font_path):
                        self._pil_font = PIL.ImageFont.truetype(self.font_path, self.font_size)
                    else:
                        # Use default font
                        self._pil_font = PIL.ImageFont.load_default()
                        # Try to get a better default font if available
                        try:
                            self._pil_font = PIL.ImageFont.truetype("Arial", self.font_size)
                        except:
                            try:
                                self._pil_font = PIL.ImageFont.truetype("DejaVuSans", self.font_size)
                            except:
                                pass  # Stick with the basic default font
                except ImportError:
                    logger.warning("PIL not available, falling back to CV2 for text rendering")
                    pass
        except Exception as e:
            logger.warning(f"Error initializing font: {e}")

    def _wrap_text(self, text: str, max_width: int) -> List[str]:
        """
        Wrap text to fit within max_width.

        Args:
            text: Text to wrap
            max_width: Maximum width in characters

        Returns:
            List of wrapped lines
        """
        if not text:
            return []

        # Split the input text into paragraphs
        paragraphs = text.split("\n")
        wrapped_lines = []

        for paragraph in paragraphs:
            if not paragraph:
                wrapped_lines.append("")
                continue

            # Estimate character width based on font size
            char_width = self.font_size * self.char_width_factor
            estimated_max_chars = max(1, int(max_width / char_width))

            # Wrap paragraph
            lines = textwrap.wrap(paragraph, width=estimated_max_chars)
            wrapped_lines.extend(lines)

        return wrapped_lines

    def _render_text_pil(
        self,
        frame: np.ndarray,
        text: str,
        position: Tuple[int, int]
    ) -> np.ndarray:
        """
        Render text using PIL.

        Args:
            frame: Frame to render text on
            text: Text to render
            position: Position to render text (x, y)

        Returns:
            Frame with text rendered
        """
        if not text or not self._pil_font:
            return frame

        # Create PIL image from frame
        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_image)

        # Calculate text size
        bbox = draw.textbbox((0, 0), text, font=self._pil_font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]

        # Adjust position based on alignment
        x, y = position
        if self.text_align == "center":
            x -= text_width // 2
        elif self.text_align == "right":
            x -= text_width

        # Draw shadow if enabled
        if self.shadow:
            shadow_x = x + self.shadow_offset
            shadow_y = y + self.shadow_offset
            draw.text((shadow_x, shadow_y), text, font=self._pil_font, fill=self.shadow_color)

        # Draw text
        draw.text((x, y), text, font=self._pil_font, fill=self.font_color)

        # Convert back to OpenCV format
        result = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

        return result

    def _render_subtitle_box(
        self,
        frame: np.ndarray,
        text_lines: List[str],
        position: str
    ) -> np.ndarray:
        """
        Render subtitle box with text.

        Args:
            frame: Frame to render on
            text_lines: Lines of text to render
            position: Position of subtitle ('top', 'bottom', 'center')

        Returns:
            Frame with subtitle box
        """
        height, width = frame.shape[:2]

        # Calculate maximum width based on frame size
        max_width = self.max_width or (width - 2 * self.margin)

        # Create PIL image for rendering
        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_image, 'RGBA')

        # Calculate text dimensions
        total_height = 0
        max_text_width = 0

        for line in text_lines:
            if not line:
                total_height += self.line_spacing
                continue

            bbox = draw.textbbox((0, 0), line, font=self._pil_font)
            line_width = bbox[2] - bbox[0]
            line_height = bbox[3] - bbox[1]

            max_text_width = max(max_text_width, line_width)
            total_height += line_height + self.line_spacing

        if total_height > 0:
            total_height -= self.line_spacing  # Remove extra spacing after last line

        # Calculate box dimensions
        box_width = max_text_width + 2 * self.padding
        box_height = total_height + 2 * self.padding

        # Calculate box position
        if position == "top":
            box_x = (width - box_width) // 2
            box_y = self.margin
        elif position == "bottom":
            box_x = (width - box_width) // 2
            box_y = height - box_height - self.margin
        else:  # "center"
            box_x = (width - box_width) // 2
            box_y = (height - box_height) // 2

        # Draw background box
        bg_color = (*self.background_color, int(255 * self.background_opacity))
        draw.rectangle(
            [(box_x, box_y), (box_x + box_width, box_y + box_height)],
            fill=bg_color
        )

        # Draw text lines
        y = box_y + self.padding
        for line in text_lines:
            if not line:
                y += self.line_spacing
                continue

            bbox = draw.textbbox((0, 0), line, font=self._pil_font)
            line_width = bbox[2] - bbox[0]
            line_height = bbox[3] - bbox[1]

            # Calculate horizontal position based on alignment
            if self.text_align == "center":
                x = box_x + (box_width - line_width) // 2
            elif self.text_align == "right":
                x = box_x + box_width - line_width - self.padding
            else:  # "left"
                x = box_x + self.padding

            # Draw shadow if enabled
            if self.shadow:
                draw.text(
                    (x + self.shadow_offset, y + self.shadow_offset),
                    line,
                    font=self._pil_font,
                    fill=self.shadow_color
                )

            # Draw text
            draw.text((x, y), line, font=self._pil_font, fill=self.font_color)

            y += line_height + self.line_spacing

        # Convert back to OpenCV format
        result = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

        return result

    def apply(self, frame: np.ndarray, text: str) -> np.ndarray:
        """
        Apply subtitle overlay to a frame.

        Args:
            frame: Video frame as numpy array
            text: Subtitle text

        Returns:
            Frame with subtitle overlay
        """
        if not text or not text.strip():
            return frame

        # Check if frame size has changed
        height, width = frame.shape[:2]
        if self._last_frame_size != (width, height):
            self._last_frame_size = (width, height)
            self._text_cache.clear()  # Clear cache when frame size changes

        # Calculate maximum width based on frame size
        max_width = self.max_width or (width - 2 * self.margin)

        # Wrap text to fit within max_width
        text_lines = self._wrap_text(text, max_width)

        # Generate cache key
        cache_key = f"{text}:{self.position}:{width}x{height}"

        # Check cache
        if cache_key in self._text_cache:
            return self._text_cache[cache_key].copy()

        # Create copy of frame
        result = frame.copy()

        # Render subtitle box with text
        result = self._render_subtitle_box(result, text_lines, self.position)

        # Cache the result (up to a reasonable size)
        if len(self._text_cache) < 100:  # Limit cache size
            self._text_cache[cache_key] = result.copy()

        return result
