"""
PyAV backend for FMUS-VID.

This module provides the PyAV implementation of the Backend interface,
focusing on low-level frame access and precise control.
"""

import av
import numpy as np
from typing import List, Dict, Tuple, Optional, Any, Union, Generator
from pathlib import Path
import logging
from fractions import Fraction
import threading
from queue import Queue
from concurrent.futures import ThreadPoolExecutor

from .backend import Backend

logger = logging.getLogger(__name__)

class PyAVBackend(Backend):
    """PyAV implementation of the Backend interface."""

    def __init__(self):
        """Initialize PyAV backend."""
        self.containers = {}  # Dictionary to store active containers
        self._thread_pool = ThreadPoolExecutor(max_workers=4)

    @staticmethod
    def is_available() -> bool:
        """Check if PyAV is available."""
        try:
            import av
            return True
        except ImportError:
            return False

    def load(self, path: Union[str, Path], **kwargs) -> Dict[str, Any]:
        """
        Load a video file using PyAV.

        Args:
            path: Path to video file
            **kwargs: Additional options
                - start: Start time in seconds
                - end: End time in seconds
                - thread_count: Number of decoding threads
                - thread_type: Threading type ("AUTO", "FRAME", "SLICE")

        Returns:
            Dictionary with video information
        """
        path = str(path)

        try:
            # Open container with options
            options = {
                "thread_count": str(kwargs.get("thread_count", "4")),
                "thread_type": kwargs.get("thread_type", "AUTO")
            }

            container = av.open(path, options=options)

            # Get video stream
            video_stream = container.streams.video[0]

            # Configure stream
            video_stream.thread_type = options["thread_type"]
            video_stream.thread_count = int(options["thread_count"])

            # Store container reference
            container_id = id(container)
            self.containers[container_id] = container

            # Calculate duration
            duration = float(video_stream.duration * video_stream.time_base)

            # Apply time range if specified
            start = kwargs.get("start", 0)
            end = kwargs.get("end")

            if start > 0 or end is not None:
                if end is None:
                    end = duration
                duration = end - start

            return {
                "container_id": container_id,
                "stream_index": video_stream.index,
                "duration": duration,
                "width": video_stream.width,
                "height": video_stream.height,
                "fps": float(video_stream.average_rate),
                "start": start,
                "end": end
            }
        except Exception as e:
            logger.error(f"Failed to load video: {e}")
            raise

    def create(self, width: int, height: int, duration: float, fps: float,
               color: Tuple[int, int, int], **kwargs) -> Dict[str, Any]:
        """
        Create a blank video using PyAV.

        Args:
            width: Frame width
            height: Frame height
            duration: Duration in seconds
            fps: Frames per second
            color: Background color (RGB)

        Returns:
            Dictionary with video information
        """
        try:
            # Create output container
            container = av.open(f"blank_{width}x{height}.mp4", mode="w")

            # Add video stream
            stream = container.add_stream("h264", rate=fps)
            stream.width = width
            stream.height = height
            stream.pix_fmt = "yuv420p"

            # Store container reference
            container_id = id(container)
            self.containers[container_id] = {
                "container": container,
                "color": color
            }

            return {
                "container_id": container_id,
                "stream_index": stream.index,
                "duration": duration,
                "width": width,
                "height": height,
                "fps": fps
            }
        except Exception as e:
            logger.error(f"Failed to create video: {e}")
            raise

    def save(self, video: Dict[str, Any], output_path: Union[str, Path],
             **kwargs) -> None:
        """
        Save video using PyAV.

        Args:
            video: Video dictionary
            output_path: Output file path
            **kwargs: Export options
                - codec: Video codec
                - pixel_format: Output pixel format
                - quality: Output quality (0-63, lower is better)
                - preset: Encoding preset
                - threads: Number of encoding threads
        """
        input_container = self.containers[video["container_id"]]
        output_path = str(output_path)

        try:
            # Open output container
            output = av.open(output_path, mode="w")

            # Add video stream
            stream = output.add_stream(
                kwargs.get("codec", "h264"),
                rate=video.get("fps", 30)
            )
            stream.width = video["width"]
            stream.height = video["height"]
            stream.pix_fmt = kwargs.get("pixel_format", "yuv420p")

            if "quality" in kwargs:
                stream.quality = kwargs["quality"]
            if "preset" in kwargs:
                stream.options["preset"] = kwargs["preset"]
            if "threads" in kwargs:
                stream.options["threads"] = str(kwargs["threads"])

            # Process frames
            for frame in self._iter_frames(video):
                packet = stream.encode(frame)
                if packet:
                    output.mux(packet)

            # Flush encoder
            packet = stream.encode(None)
            if packet:
                output.mux(packet)

            # Close containers
            output.close()
            if isinstance(input_container, av.container.Container):
                input_container.close()

            # Clean up
            del self.containers[video["container_id"]]

        except Exception as e:
            logger.error(f"Failed to save video: {e}")
            raise

    def _iter_frames(self, video: Dict[str, Any]) -> Generator[av.VideoFrame, None, None]:
        """Iterate over video frames."""
        container = self.containers[video["container_id"]]

        if isinstance(container, dict):  # Blank video
            # Generate blank frames
            color = container["color"]
            frame = av.VideoFrame(video["width"], video["height"], "rgb24")
            frame.planes[0].update(
                np.full((video["height"], video["width"], 3), color, dtype=np.uint8)
            )
            num_frames = int(video["duration"] * video["fps"])
            for _ in range(num_frames):
                yield frame
        else:  # Normal video
            stream = container.streams.video[0]
            stream.thread_type = "AUTO"
            stream.thread_count = 4

            start_time = video.get("start", 0)
            end_time = video.get("end")

            if start_time > 0:
                container.seek(int(start_time * stream.time_base.denominator))

            for frame in container.decode(video=0):
                if end_time and frame.time > end_time:
                    break
                yield frame

    def extract_frame(self, video: Dict[str, Any], time: float) -> np.ndarray:
        """
        Extract a single frame at specified time.

        Args:
            video: Video dictionary
            time: Time in seconds

        Returns:
            Frame as numpy array (RGB)
        """
        container = self.containers[video["container_id"]]
        if isinstance(container, dict):
            # Return blank frame
            return np.full(
                (video["height"], video["width"], 3),
                container["color"],
                dtype=np.uint8
            )

        stream = container.streams.video[0]
        pts = int(time / stream.time_base)
        container.seek(pts, stream=stream)

        for frame in container.decode(video=0):
            if frame.pts >= pts:
                return frame.to_ndarray(format="rgb24")

        raise ValueError(f"Frame at time {time} not found")

    def get_frame_batch(self, video: Dict[str, Any],
                       times: List[float]) -> List[np.ndarray]:
        """
        Extract multiple frames at specified times.

        Args:
            video: Video dictionary
            times: List of times in seconds

        Returns:
            List of frames as numpy arrays (RGB)
        """
        frames = []
        futures = []

        # Submit frame extraction tasks
        for time in times:
            future = self._thread_pool.submit(self.extract_frame, video, time)
            futures.append(future)

        # Collect results
        for future in futures:
            frames.append(future.result())

        return frames

    def decode_stream(self, video: Dict[str, Any],
                     frame_callback: Callable[[np.ndarray, float], None],
                     batch_size: int = 10) -> None:
        """
        Decode video stream with frame callback.

        Args:
            video: Video dictionary
            frame_callback: Function to call with (frame, timestamp)
            batch_size: Number of frames to process in parallel
        """
        frames = []
        timestamps = []

        for frame in self._iter_frames(video):
            frames.append(frame.to_ndarray(format="rgb24"))
            timestamps.append(float(frame.pts * frame.time_base))

            if len(frames) >= batch_size:
                # Process batch
                for f, t in zip(frames, timestamps):
                    frame_callback(f, t)
                frames.clear()
                timestamps.clear()

        # Process remaining frames
        for f, t in zip(frames, timestamps):
            frame_callback(f, t)

    def cleanup(self):
        """Clean up all containers."""
        for container in self.containers.values():
            try:
                if isinstance(container, av.container.Container):
                    container.close()
            except:
                pass
        self.containers.clear()
        self._thread_pool.shutdown()

    def __del__(self):
        """Ensure cleanup on deletion."""
        self.cleanup()

    def grid(self, videos: List[Any], rows: int, cols: int) -> Any:
        """Arrange videos in a grid layout."""
        # This is a more complex operation that would require
        # frame-by-frame processing and composition
        # For simplicity, just log a warning
        logger.warning("Grid operation not yet implemented for PyAV backend")
        return videos[0] if videos else None

    def add_subtitles(self, video: Any, entries: List,
                     font: str = "Arial", size: int = 24,
                     color: Union[str, Tuple[int, int, int]] = "white",
                     position: Optional[Tuple[int, int]] = None) -> Any:
        """
        Add subtitles from a list of entries.

        Args:
            video: Video dictionary
            entries: List of SubtitleEntry objects
            font: Font name or path
            size: Font size in pixels
            color: Text color (name or RGB tuple)
            position: Default position (None for bottom center)

        Returns:
            Updated video dictionary
        """
        # This would require frame-by-frame processing with PIL
        # For simplicity, just log a warning
        logger.warning("Subtitle operation not yet implemented for PyAV backend")
        return video

    def add_subtitle_text(self, video: Any, entry,
                         font: str = "Arial", size: int = 24,
                         color: Union[str, Tuple[int, int, int]] = "white") -> Any:
        """
        Add a single subtitle entry.

        Args:
            video: Video dictionary
            entry: SubtitleEntry object
            font: Font name or path
            size: Font size in pixels
            color: Text color (name or RGB tuple)

        Returns:
            Updated video dictionary
        """
        # This would require frame-by-frame processing with PIL
        # For simplicity, just log a warning
        logger.warning("Subtitle operation not yet implemented for PyAV backend")
        return video
