"""
Real-time Video Player with Collision Prediction Overlay and LED Feedback

This module provides real-time video playback with:
- GPU-accelerated collision prediction via BADAS
- Visual overlays showing risk levels and probability bars
- Physical LED feedback (optional)
- Headless mode for processing without display

Usage:
    # As a module
    from badas.inference.realtime import RealtimeVideoPlayer

    player = RealtimeVideoPlayer(
        model_path="/path/to/model.ckpt",
        led_port='auto',  # or None to disable
    )
    player.play_video("video.mp4")
    player.close()

    # From command line
    python -m badas.inference.realtime video.mp4 --model model.ckpt --led-port auto

Controls during playback:
    q / ESC: Quit
    SPACE: Pause/Resume
    r: Restart video
"""

import cv2
import numpy as np
import time
import argparse
from pathlib import Path
from typing import Optional, Dict, Any

from badas.inference.led import (
    BaseLEDController,
    DummyLEDController,
    create_led_controller,
)
from badas.visualization.overlay import OverlayRenderer


class RealtimeVideoPlayer:
    """
    Real-time video player with inference and overlay.

    Supports:
    - PyTorch, ONNX, and TensorRT backends via BADAS
    - Physical LED control (optional)
    - Headless mode (process without display)
    - Video output saving

    Args:
        model_path: Path to model file (.ckpt, .trt, .onnx)
        led_port: Serial port for LED ('auto', specific port, or None to disable)
        target_fps: Model's expected FPS (default 8)
        smoothing_alpha: EMA smoothing factor (0=smooth, 1=responsive)
        headless: If True, skip display and only process/save
        quiet: If True, suppress output messages
    """

    def __init__(
        self,
        model_path: str,
        led_port: Optional[str] = 'auto',
        target_fps: int = 8,
        smoothing_alpha: float = 0.3,
        headless: bool = False,
        quiet: bool = False,
        **badas_kwargs
    ):
        from badas.inference import BADAS, BADASConfig, SmoothingConfig

        self.headless = headless
        self.target_fps = target_fps
        self.smoothing_alpha = smoothing_alpha
        self.quiet = quiet

        # Determine model info for display
        model_ext = Path(model_path).suffix.lower()
        if model_ext in ('.trt', '.engine'):
            self.model_info = f"TensorRT @ {target_fps}fps"
        elif model_ext == '.onnx':
            self.model_info = f"ONNX @ {target_fps}fps"
        else:
            self.model_info = f"PyTorch @ {target_fps}fps"

        self._log(f"Loading model from: {model_path}")

        # Create BADAS predictor - uses new repo's inference
        config = BADASConfig(
            use_compile=False,  # Faster startup
            smoothing=SmoothingConfig(enabled=False),  # We handle smoothing ourselves
            quiet=True,
            **badas_kwargs
        )

        self.predictor = BADAS(model_path, config=config)
        self._log("Model loaded!")

        # LED controller
        if led_port is not None:
            self.led = create_led_controller(led_port, fallback_to_dummy=True)
            if isinstance(self.led, DummyLEDController):
                self._log("Using simulated LED indicator")
            else:
                self._log(f"LED controller connected on {self.led.port}")
        else:
            self.led = DummyLEDController()
            self._log("LED disabled")

        # Overlay renderer
        self.renderer = OverlayRenderer()

        # State
        self._smoothed_prob = 0.0
        self._is_ready = False
        self._paused = False
        self._buffer_fill = 0

    def _log(self, msg: str):
        if not self.quiet:
            print(msg)

    def reset(self):
        """Reset state for new video."""
        self.predictor.reset_streaming_state()
        self._smoothed_prob = 0.0
        self._is_ready = False
        self._buffer_fill = 0
        self.led.blue()

    def process_frame(self, frame: np.ndarray) -> Optional[Dict[str, Any]]:
        """
        Process a frame and return prediction info.

        Args:
            frame: BGR frame from video/camera

        Returns:
            Prediction dict or None if still buffering
        """
        start = time.perf_counter()

        result = self.predictor.predict_frame(frame, apply_smoothing=False)

        if result.get('status') == 'buffering':
            self._buffer_fill = result.get('buffer_fill', 0)
            return None

        if not self._is_ready:
            self._is_ready = True

        # Get raw probability
        raw_prob = result.get('raw_probability', result.get('probability', 0.0))

        # Apply EMA smoothing
        if self._smoothed_prob == 0.0:
            self._smoothed_prob = raw_prob
        else:
            self._smoothed_prob = (
                self.smoothing_alpha * raw_prob +
                (1 - self.smoothing_alpha) * self._smoothed_prob
            )

        # Update LED
        self.led.set_risk(self._smoothed_prob)

        latency = (time.perf_counter() - start) * 1000

        return {
            'probability': self._smoothed_prob,
            'raw_probability': raw_prob,
            'latency_ms': latency,
        }

    def play_video(
        self,
        video_path: str,
        window_name: str = "BADAS - Collision Prediction",
        output_path: Optional[str] = None,
        show_player: bool = True
    ) -> None:
        """
        Play video with real-time inference and overlay.

        Args:
            video_path: Path to video file
            window_name: Name of the display window
            output_path: If provided, save output video to this path
            show_player: If True, show video player window (ignored if headless)
        """
        self.reset()

        # Override headless if show_player is False
        headless = self.headless or not show_player

        # Open video
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Could not open video: {video_path}")

        video_fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        duration = total_frames / video_fps

        self._log(f"\nVideo: {video_path}")
        self._log(f"Resolution: {width}x{height} @ {video_fps:.0f}fps")
        self._log(f"Duration: {duration:.1f}s ({total_frames} frames)")
        self._log(f"Model sampling: {self.target_fps}fps")

        # Setup video writer if saving
        video_writer = None
        if output_path:
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            video_writer = cv2.VideoWriter(output_path, fourcc, video_fps, (width, height))
            self._log(f"Saving output to: {output_path}")

        # Calculate frame interval for model sampling
        frame_interval = video_fps / self.target_fps

        # Create window only if not headless
        if not headless:
            self._log("\nControls:")
            self._log("  q/ESC: Quit")
            self._log("  SPACE: Pause/Resume")
            self._log("  r: Restart")
            self._log("-" * 50)
            try:
                cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
                cv2.resizeWindow(window_name, min(width, 1280), min(height, 720))
            except cv2.error as e:
                self._log(f"\nWarning: Cannot create display window: {e}")
                self._log("Switching to headless mode (LED + save only)")
                headless = True
        else:
            self._log("\nRunning in headless mode (no display)")
            self._log("-" * 50)

        current_frame = 0
        next_sample = 0.0
        last_result = None
        frame_time = 1.0 / video_fps
        start_time = time.perf_counter()

        try:
            while True:
                if not self._paused:
                    loop_start = time.perf_counter()

                    ret, frame = cap.read()
                    if not ret:
                        if not headless:
                            self._log("\nVideo ended. Press 'r' to restart or 'q' to quit.")
                            self._paused = True
                            continue
                        else:
                            self._log("\nVideo processing complete!")
                            break

                    current_time = current_frame / video_fps

                    # Process frame for inference at target FPS
                    if current_frame >= int(next_sample):
                        result = self.process_frame(frame)
                        if result:
                            last_result = result
                            # Print progress to terminal
                            prob = result['probability']
                            raw = result['raw_probability']
                            indicator = '\U0001F534' if prob >= 0.8 else '\U0001F7E0' if prob >= 0.5 else '\U0001F7E2'
                            self._log(f"{indicator} t={current_time:5.2f}s | Risk: {prob:5.1%} (raw: {raw:5.1%}) | Latency: {result['latency_ms']:.0f}ms")
                        next_sample += frame_interval

                    # Draw overlay
                    if current_time < 2.0:
                        # Still initializing
                        display_frame = self.renderer.draw_initialization_overlay(
                            frame, current_time,
                            buffer_fill=self._buffer_fill,
                            buffer_size=self.predictor.window_size
                        )
                    elif last_result:
                        # Have prediction
                        display_frame = self.renderer.draw_prediction_overlay(
                            frame,
                            last_result['probability'],
                            last_result['raw_probability'],
                            current_time,
                            last_result['latency_ms'],
                            self.led.state,
                            self.model_info
                        )
                    else:
                        display_frame = frame

                    # Save frame if writing
                    if video_writer:
                        video_writer.write(display_frame)

                    # Show frame if not headless
                    if not headless:
                        cv2.imshow(window_name, display_frame)

                    current_frame += 1

                    # Control playback speed (only for display mode)
                    if not headless:
                        elapsed = time.perf_counter() - loop_start
                        wait_time = max(1, int((frame_time - elapsed) * 1000))
                    else:
                        wait_time = 1  # Minimal wait in headless mode
                else:
                    wait_time = 50  # Check input every 50ms when paused

                # Handle keyboard input (only if display is available)
                if not headless:
                    key = cv2.waitKey(wait_time) & 0xFF

                    if key == ord('q') or key == 27:  # q or ESC
                        self._log("\nQuitting...")
                        break
                    elif key == ord(' '):  # SPACE
                        self._paused = not self._paused
                        self._log("PAUSED" if self._paused else "RESUMED")
                    elif key == ord('r'):  # r
                        self._log("\nRestarting video...")
                        self.reset()
                        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                        current_frame = 0
                        next_sample = 0.0
                        last_result = None
                        self._paused = False

        finally:
            cap.release()
            if video_writer:
                video_writer.release()
                self._log(f"\nSaved output video to: {output_path}")
            if not headless:
                cv2.destroyAllWindows()
            self.led.blue()

            # Print summary
            elapsed_total = time.perf_counter() - start_time
            self._log(f"Processed {current_frame} frames in {elapsed_total:.1f}s")

    def close(self):
        """Clean up resources."""
        self.led.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False


def main():
    """Command-line interface for realtime video player."""
    parser = argparse.ArgumentParser(
        description="Real-time video player with BADAS collision prediction overlay",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Controls during playback:
  q / ESC   - Quit
  SPACE     - Pause/Resume
  r         - Restart video

Risk levels:
  Green  (< 50%)  - Normal / Low risk
  Orange (50-80%) - Caution / Medium risk
  Red    (>= 80%) - Danger / High risk

Examples:
  # Play video with default settings
  python -m badas.inference.realtime video.mp4 --model model.ckpt

  # Save output without display
  python -m badas.inference.realtime video.mp4 --model model.ckpt --save output.mp4 --headless

  # Disable LED
  python -m badas.inference.realtime video.mp4 --model model.ckpt --no-led
        """
    )
    parser.add_argument(
        'video',
        help='Path to video file'
    )
    parser.add_argument(
        '--model', '-m',
        required=True,
        help='Path to model file (.ckpt, .trt, or .onnx)'
    )
    parser.add_argument(
        '--led-port', '-l',
        default='auto',
        help='Serial port for LED controller (default: auto-detect)'
    )
    parser.add_argument(
        '--no-led', action='store_true',
        help='Disable LED controller (use simulated indicator)'
    )
    parser.add_argument(
        '--fps', '-f', type=int, default=8,
        help='Target FPS for model inference (default: 8)'
    )
    parser.add_argument(
        '--smoothing', '-s', type=float, default=0.3,
        help='EMA smoothing factor 0-1 (default: 0.3, lower=smoother)'
    )
    parser.add_argument(
        '--save', '-o',
        help='Save output video to this path'
    )
    parser.add_argument(
        '--headless', action='store_true',
        help='Run without display (requires --save to be useful)'
    )
    parser.add_argument(
        '--no-player', action='store_true',
        help='Disable video player window (same as --headless)'
    )
    parser.add_argument(
        '--quiet', '-q', action='store_true',
        help='Suppress output messages'
    )

    args = parser.parse_args()

    # Determine LED port
    led_port = None if args.no_led else args.led_port

    # Auto-enable headless if no-player specified
    headless = args.headless or args.no_player

    if not args.quiet:
        print("=" * 50)
        print("BADAS - Real-time Collision Prediction")
        print("=" * 50)

    try:
        with RealtimeVideoPlayer(
            model_path=args.model,
            led_port=led_port,
            target_fps=args.fps,
            smoothing_alpha=args.smoothing,
            headless=headless,
            quiet=args.quiet
        ) as player:
            player.play_video(
                args.video,
                output_path=args.save,
                show_player=not args.no_player
            )

    except KeyboardInterrupt:
        print("\nInterrupted by user")
    except Exception as e:
        print(f"\nError: {e}")
        raise


if __name__ == '__main__':
    main()
