"""
Streaming functionality for FMUS-VID.

This module provides classes for handling video streams, including reading from
cameras, network streams, and writing to streaming outputs.
"""

from typing import Union, Optional, Dict, Any, Iterator, Tuple
from pathlib import Path
import logging
import time
from abc import ABC, abstractmethod
from queue import Queue
from threading import Thread, Event
import cv2
import numpy as np
from PIL import Image
import subprocess
import tempfile
import os

logger = logging.getLogger(__name__)

class VideoStream(ABC):
    """Abstract base class for video streams."""

    def __init__(self):
        self._is_running = False
        self._stop_event = Event()
        self._frame_queue = Queue(maxsize=30)  # Buffer 30 frames
        self._thread: Optional[Thread] = None

    @abstractmethod
    def _read_frame(self) -> Optional[np.ndarray]:
        """Read a single frame from the stream."""
        pass

    def start(self):
        """Start the video stream."""
        if self._is_running:
            return

        self._is_running = True
        self._stop_event.clear()
        self._thread = Thread(target=self._stream_worker, daemon=True)
        self._thread.start()

    def stop(self):
        """Stop the video stream."""
        self._stop_event.set()
        if self._thread:
            self._thread.join()
        self._is_running = False
        self._frame_queue = Queue(maxsize=30)

    def _stream_worker(self):
        """Worker thread for reading frames."""
        while not self._stop_event.is_set():
            frame = self._read_frame()
            if frame is None:
                break

            # Convert to PIL Image
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_frame = Image.fromarray(frame_rgb)

            try:
                self._frame_queue.put(pil_frame, timeout=1)
            except Queue.Full:
                # Drop oldest frame if queue is full
                try:
                    self._frame_queue.get_nowait()
                    self._frame_queue.put(pil_frame)
                except Queue.Empty:
                    pass

    def read(self, timeout: float = 1.0) -> Optional[Image.Image]:
        """
        Read a frame from the stream.

        Args:
            timeout: Time to wait for frame in seconds

        Returns:
            PIL Image or None if no frame available
        """
        if not self._is_running:
            self.start()

        try:
            return self._frame_queue.get(timeout=timeout)
        except Queue.Empty:
            return None

    def __iter__(self) -> Iterator[Image.Image]:
        """Iterate over frames in the stream."""
        if not self._is_running:
            self.start()

        while self._is_running:
            frame = self.read()
            if frame is not None:
                yield frame

    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.stop()

class CameraStream(VideoStream):
    """Stream from a camera device."""

    def __init__(self, device_id: int = 0, width: Optional[int] = None,
                 height: Optional[int] = None, fps: Optional[int] = None):
        """
        Initialize camera stream.

        Args:
            device_id: Camera device ID
            width: Desired frame width
            height: Desired frame height
            fps: Desired frames per second
        """
        super().__init__()
        self.device_id = device_id
        self.width = width
        self.height = height
        self.fps = fps
        self._cap = None

    def start(self):
        """Start the camera stream."""
        self._cap = cv2.VideoCapture(self.device_id)

        if self.width and self.height:
            self._cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
            self._cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)

        if self.fps:
            self._cap.set(cv2.CAP_PROP_FPS, self.fps)

        if not self._cap.isOpened():
            raise RuntimeError(f"Failed to open camera {self.device_id}")

        super().start()

    def stop(self):
        """Stop the camera stream."""
        super().stop()
        if self._cap:
            self._cap.release()
            self._cap = None

    def _read_frame(self) -> Optional[np.ndarray]:
        """Read a frame from the camera."""
        if not self._cap or not self._cap.isOpened():
            return None

        ret, frame = self._cap.read()
        return frame if ret else None

class NetworkStream(VideoStream):
    """Stream from a network source (RTSP, HTTP, etc.)."""

    def __init__(self, url: str, buffer_size: int = 30):
        """
        Initialize network stream.

        Args:
            url: Stream URL (e.g., RTSP or HTTP)
            buffer_size: Frame buffer size
        """
        super().__init__()
        self.url = url
        self._frame_queue = Queue(maxsize=buffer_size)
        self._cap = None

    def start(self):
        """Start the network stream."""
        self._cap = cv2.VideoCapture(self.url)

        if not self._cap.isOpened():
            raise RuntimeError(f"Failed to open stream: {self.url}")

        super().start()

    def stop(self):
        """Stop the network stream."""
        super().stop()
        if self._cap:
            self._cap.release()
            self._cap = None

    def _read_frame(self) -> Optional[np.ndarray]:
        """Read a frame from the network stream."""
        if not self._cap or not self._cap.isOpened():
            return None

        ret, frame = self._cap.read()
        return frame if ret else None

class RTMPStream:
    """Stream video to RTMP server (YouTube, Twitch, etc.)."""

    def __init__(self, rtmp_url: str, width: int, height: int,
                 fps: float = 30.0, bitrate: str = '2500k',
                 video_codec: str = 'libx264', audio_codec: str = 'aac',
                 audio_bitrate: str = '128k', preset: str = 'medium',
                 format: str = 'flv'):
        """
        Initialize RTMP stream.

        Args:
            rtmp_url: RTMP server URL
            width: Video width
            height: Video height
            fps: Frames per second
            bitrate: Video bitrate
            video_codec: Video codec
            audio_codec: Audio codec
            audio_bitrate: Audio bitrate
            preset: Encoding preset (ultrafast, superfast, veryfast, faster, fast, medium, slow, slower, veryslow)
            format: Output format (usually flv for RTMP)
        """
        self.rtmp_url = rtmp_url
        self.width = width
        self.height = height
        self.fps = fps
        self.bitrate = bitrate
        self.video_codec = video_codec
        self.audio_codec = audio_codec
        self.audio_bitrate = audio_bitrate
        self.preset = preset
        self.format = format

        self._ffmpeg_process = None
        self._temp_dir = None
        self._is_running = False

    def start(self):
        """Start the RTMP stream."""
        if self._is_running:
            return

        # Create temporary directory for any needed files
        self._temp_dir = tempfile.mkdtemp(prefix="fmusvid_rtmp_")

        # Build FFmpeg command
        command = [
            'ffmpeg',
            # Input options
            '-f', 'rawvideo',           # Input format
            '-pix_fmt', 'rgb24',        # Pixel format
            '-s', f'{self.width}x{self.height}',  # Frame size
            '-r', str(self.fps),        # Frame rate
            '-i', 'pipe:',              # Read from stdin

            # Audio input (silent)
            '-f', 'lavfi',              # Use libavfilter
            '-i', 'anullsrc',           # Generate silent audio

            # Output options
            '-c:v', self.video_codec,   # Video codec
            '-preset', self.preset,     # Encoding preset
            '-b:v', self.bitrate,       # Video bitrate
            '-c:a', self.audio_codec,   # Audio codec
            '-b:a', self.audio_bitrate, # Audio bitrate
            '-f', self.format,          # Output format
            '-flvflags', 'no_duration_filesize',  # FLV flags
            '-g', str(int(self.fps * 2)),  # GOP size (2 seconds)
            '-bufsize', '2M',           # Buffer size
            '-maxrate', '2500k',        # Max bitrate

            # RTMP URL
            self.rtmp_url
        ]

        # Start FFmpeg process
        self._ffmpeg_process = subprocess.Popen(
            command,
            stdin=subprocess.PIPE,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            bufsize=10**8  # Large buffer
        )

        self._is_running = True
        logger.info(f"Started RTMP stream to {self.rtmp_url}")

    def write(self, frame: Union[np.ndarray, Image.Image]):
        """
        Write a frame to the RTMP stream.

        Args:
            frame: Frame to write (numpy array or PIL Image)
        """
        if not self._is_running:
            self.start()

        # Convert frame to numpy array if it's a PIL Image
        if isinstance(frame, Image.Image):
            frame = np.array(frame)

        # Convert BGR to RGB if needed
        if frame.shape[2] == 3 and frame.dtype == np.uint8:
            if not isinstance(frame, np.ndarray):
                frame = np.array(frame)
            # OpenCV uses BGR, convert to RGB for ffmpeg
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # Resize frame if needed
        if frame.shape[0] != self.height or frame.shape[1] != self.width:
            frame = cv2.resize(frame, (self.width, self.height))

        # Ensure the process is still running
        if self._ffmpeg_process and self._ffmpeg_process.poll() is None:
            try:
                # Write frame to FFmpeg's stdin
                self._ffmpeg_process.stdin.write(frame.tobytes())
                self._ffmpeg_process.stdin.flush()
            except (BrokenPipeError, IOError) as e:
                logger.error(f"Error writing to RTMP stream: {e}")
                self.stop()
                # Try to restart
                self.start()
        else:
            logger.warning("RTMP stream process not running, attempting to restart")
            self.stop()
            self.start()

    def stop(self):
        """Stop the RTMP stream."""
        if not self._is_running:
            return

        # Close FFmpeg process
        if self._ffmpeg_process:
            try:
                self._ffmpeg_process.stdin.close()
                self._ffmpeg_process.wait(timeout=5)
            except (subprocess.TimeoutExpired, BrokenPipeError, IOError):
                self._ffmpeg_process.kill()

            self._ffmpeg_process = None

        # Clean up temporary directory
        if self._temp_dir and os.path.exists(self._temp_dir):
            try:
                import shutil
                shutil.rmtree(self._temp_dir)
            except OSError as e:
                logger.warning(f"Failed to remove temp directory: {e}")

        self._is_running = False
        logger.info("Stopped RTMP stream")

    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.stop()

class StreamWriter:
    """Class for writing video streams."""

    def __init__(self, output: Union[str, Path], width: int, height: int,
                 fps: float = 30.0, codec: str = 'avc1', bitrate: str = '2000k'):
        """
        Initialize stream writer.

        Args:
            output: Output path or URL
            width: Frame width
            height: Frame height
            fps: Frames per second
            codec: Video codec (e.g., 'avc1', 'vp8')
            bitrate: Video bitrate
        """
        self.output = str(output)
        self.width = width
        self.height = height
        self.fps = fps
        self.codec = codec
        self.bitrate = bitrate

        self._writer = None
        self._start_time = None
        self._frame_count = 0

    def start(self):
        """Start the stream writer."""
        fourcc = cv2.VideoWriter_fourcc(*self.codec)
        self._writer = cv2.VideoWriter(
            self.output, fourcc, self.fps,
            (self.width, self.height)
        )

        if not self._writer.isOpened():
            raise RuntimeError(f"Failed to create video writer: {self.output}")

        self._start_time = time.time()
        self._frame_count = 0

    def write(self, frame: Union[np.ndarray, Image.Image]):
        """
        Write a frame to the stream.

        Args:
            frame: Frame to write (numpy array or PIL Image)
        """
        if not self._writer:
            self.start()

        # Convert PIL Image to numpy array if needed
        if isinstance(frame, Image.Image):
            frame = cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR)

        self._writer.write(frame)
        self._frame_count += 1

        # Maintain correct frame rate
        expected_time = self._frame_count / self.fps
        actual_time = time.time() - self._start_time
        if actual_time < expected_time:
            time.sleep(expected_time - actual_time)

    def stop(self):
        """Stop the stream writer."""
        if self._writer:
            self._writer.release()
            self._writer = None

    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.stop()
