#!/usr/bin/env python3
"""
Drone Camera Viewer - Real-time camera feed display.

Displays the drone's camera feed from Gazebo simulation in a separate window.
Works with ROS 2 image topics.

Usage:
    python camera_viewer.py                     # Default topic /drone/camera
    python camera_viewer.py --topic /camera/image_raw
    python camera_viewer.py --fps 60            # Higher framerate
    python camera_viewer.py --record output.avi # Record to file

Topics:
    - /drone/camera (Gazebo simulation)
    - /camera/image_raw (generic)
    - /ap/camera (ArduPilot if available)
"""

import argparse
import sys
import time
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from sensor_msgs.msg import Image

# Try to import OpenCV
try:
    import cv2
    import numpy as np
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    print("Warning: OpenCV not installed. Install with: pip install opencv-python")

# Try to import cv_bridge (optional, for more robust conversion)
try:
    from cv_bridge import CvBridge
    CV_BRIDGE_AVAILABLE = True
except ImportError:
    CV_BRIDGE_AVAILABLE = False


class CameraViewer(Node):
    """ROS 2 node that displays camera feed in a window."""
    
    def __init__(
        self,
        topic: str = '/drone/camera',
        fps: int = 30,
        window_name: str = 'Drone Camera',
        record_file: Optional[str] = None,
        show_info: bool = True
    ):
        super().__init__('camera_viewer')
        
        self._topic = topic
        self._target_fps = fps
        self._window_name = window_name
        self._show_info = show_info
        self._record_file = record_file
        
        # State
        self._latest_frame: Optional[np.ndarray] = None
        self._frame_count = 0
        self._fps_actual = 0.0
        self._last_fps_time = time.time()
        self._fps_frame_count = 0
        self._running = True
        
        # Video writer for recording
        self._video_writer: Optional[cv2.VideoWriter] = None
        
        # CV Bridge for image conversion
        if CV_BRIDGE_AVAILABLE:
            self._bridge = CvBridge()
        else:
            self._bridge = None
        
        # QoS for camera topics (best effort for performance)
        camera_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1
        )
        
        # Subscribe to camera topic
        self._camera_sub = self.create_subscription(
            Image,
            topic,
            self._camera_callback,
            camera_qos
        )
        
        self.get_logger().info('=' * 50)
        self.get_logger().info('Camera Viewer Starting...')
        self.get_logger().info(f'  Topic: {topic}')
        self.get_logger().info(f'  Target FPS: {fps}')
        if record_file:
            self.get_logger().info(f'  Recording to: {record_file}')
        self.get_logger().info('=' * 50)
        self.get_logger().info('Press Q or ESC in the window to quit')
        
        # Create display timer at target FPS
        display_period = 1.0 / fps
        self._display_timer = self.create_timer(display_period, self._display_frame)
    
    def _camera_callback(self, msg: Image):
        """Process incoming camera image."""
        try:
            # Convert ROS Image to OpenCV format
            if self._bridge and CV_BRIDGE_AVAILABLE:
                # Use cv_bridge for robust conversion
                self._latest_frame = self._bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            else:
                # Manual conversion
                self._latest_frame = self._ros_image_to_cv2(msg)
            
            self._frame_count += 1
            self._fps_frame_count += 1
            
        except Exception as e:
            self.get_logger().warning(f'Failed to convert image: {e}')
    
    def _ros_image_to_cv2(self, msg: Image) -> np.ndarray:
        """Convert ROS Image message to OpenCV format."""
        # Get image dimensions
        height = msg.height
        width = msg.width
        encoding = msg.encoding
        
        # Convert based on encoding
        if encoding in ['rgb8', 'bgr8']:
            # 3-channel color image
            channels = 3
            dtype = np.uint8
        elif encoding in ['rgba8', 'bgra8']:
            # 4-channel color image
            channels = 4
            dtype = np.uint8
        elif encoding in ['mono8', '8UC1']:
            # Grayscale
            channels = 1
            dtype = np.uint8
        elif encoding in ['mono16', '16UC1']:
            channels = 1
            dtype = np.uint16
        elif encoding == '32FC1':
            channels = 1
            dtype = np.float32
        else:
            # Try to handle as 3-channel uint8
            channels = 3
            dtype = np.uint8
        
        # Reshape data to image
        if channels == 1:
            frame = np.frombuffer(msg.data, dtype=dtype).reshape((height, width))
            # Convert grayscale to BGR for display
            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
        else:
            frame = np.frombuffer(msg.data, dtype=dtype).reshape((height, width, channels))
            
            # Convert to BGR if needed
            if encoding == 'rgb8':
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            elif encoding == 'rgba8':
                frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)
            elif encoding == 'bgra8':
                frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
        
        return frame
    
    def _display_frame(self):
        """Display the latest frame."""
        if not self._running:
            return
        
        # Calculate actual FPS
        current_time = time.time()
        elapsed = current_time - self._last_fps_time
        if elapsed >= 1.0:
            self._fps_actual = self._fps_frame_count / elapsed
            self._fps_frame_count = 0
            self._last_fps_time = current_time
        
        # Check for window close or key press
        key = cv2.waitKey(1) & 0xFF
        if key in [ord('q'), ord('Q'), 27]:  # Q or ESC
            self._running = False
            self.get_logger().info('Shutting down...')
            rclpy.shutdown()
            return
        
        # Display frame if available
        if self._latest_frame is not None:
            display_frame = self._latest_frame.copy()
            
            # Add info overlay
            if self._show_info:
                self._add_info_overlay(display_frame)
            
            # Show frame
            cv2.imshow(self._window_name, display_frame)
            
            # Record if enabled
            if self._record_file and self._video_writer is None:
                # Initialize video writer with frame dimensions
                h, w = display_frame.shape[:2]
                fourcc = cv2.VideoWriter_fourcc(*'XVID')
                self._video_writer = cv2.VideoWriter(
                    self._record_file, fourcc, self._target_fps, (w, h)
                )
            
            if self._video_writer:
                self._video_writer.write(display_frame)
        else:
            # No frame yet - show waiting message
            waiting_frame = np.zeros((480, 640, 3), dtype=np.uint8)
            cv2.putText(
                waiting_frame,
                'Waiting for camera feed...',
                (120, 240),
                cv2.FONT_HERSHEY_SIMPLEX,
                1.0,
                (255, 255, 255),
                2
            )
            cv2.putText(
                waiting_frame,
                f'Topic: {self._topic}',
                (100, 280),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7,
                (128, 128, 128),
                1
            )
            cv2.imshow(self._window_name, waiting_frame)
    
    def _add_info_overlay(self, frame: np.ndarray):
        """Add information overlay to frame."""
        h, w = frame.shape[:2]
        
        # Semi-transparent background for text
        overlay = frame.copy()
        cv2.rectangle(overlay, (5, 5), (200, 85), (0, 0, 0), -1)
        cv2.addWeighted(overlay, 0.5, frame, 0.5, 0, frame)
        
        # FPS display
        cv2.putText(
            frame,
            f'FPS: {self._fps_actual:.1f}',
            (10, 25),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (0, 255, 0),
            2
        )
        
        # Frame count
        cv2.putText(
            frame,
            f'Frame: {self._frame_count}',
            (10, 50),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (0, 255, 0),
            2
        )
        
        # Resolution
        cv2.putText(
            frame,
            f'Size: {w}x{h}',
            (10, 75),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (0, 255, 0),
            2
        )
        
        # Recording indicator
        if self._record_file:
            cv2.circle(frame, (w - 20, 20), 10, (0, 0, 255), -1)
            cv2.putText(
                frame,
                'REC',
                (w - 60, 25),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 0, 255),
                2
            )
    
    def shutdown(self):
        """Clean up resources."""
        self._running = False
        if self._video_writer:
            self._video_writer.release()
        cv2.destroyAllWindows()


def find_camera_topics() -> list:
    """Find available camera topics."""
    import subprocess
    try:
        result = subprocess.run(
            ['ros2', 'topic', 'list'],
            capture_output=True,
            text=True,
            timeout=5.0
        )
        topics = result.stdout.strip().split('\n')
        
        # Filter for likely camera topics
        camera_keywords = ['camera', 'image', 'rgb', 'depth']
        camera_topics = [
            t for t in topics
            if any(kw in t.lower() for kw in camera_keywords)
        ]
        return camera_topics
    except Exception:
        return []


def main():
    if not CV2_AVAILABLE:
        print("ERROR: OpenCV is required for camera viewer")
        print("Install with: pip install opencv-python")
        sys.exit(1)
    
    parser = argparse.ArgumentParser(
        description='Drone Camera Viewer - Display camera feed in real-time',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python camera_viewer.py                        # Default topic
  python camera_viewer.py --topic /camera/image_raw
  python camera_viewer.py --fps 60               # Higher framerate
  python camera_viewer.py --record flight.avi    # Record video
  python camera_viewer.py --list                 # List camera topics

Common topics:
  /drone/camera          - Gazebo simulation
  /camera/image_raw      - Generic camera
  /camera/color/image_raw - RGB camera
  /camera/depth/image_raw - Depth camera
        """
    )
    
    parser.add_argument(
        '--topic', '-t', type=str, default='/drone/camera',
        help='Camera image topic (default: /drone/camera)'
    )
    parser.add_argument(
        '--fps', type=int, default=30,
        help='Target display framerate (default: 30)'
    )
    parser.add_argument(
        '--record', '-r', type=str, default=None,
        help='Record video to file (e.g., output.avi)'
    )
    parser.add_argument(
        '--no-info', action='store_true',
        help='Hide info overlay'
    )
    parser.add_argument(
        '--list', '-l', action='store_true',
        help='List available camera topics and exit'
    )
    parser.add_argument(
        '--window-name', '-w', type=str, default='Drone Camera',
        help='Window title'
    )
    
    args = parser.parse_args()
    
    # List topics and exit
    if args.list:
        print("Searching for camera topics...")
        topics = find_camera_topics()
        if topics:
            print("\nAvailable camera topics:")
            for t in topics:
                print(f"  {t}")
        else:
            print("\nNo camera topics found. Is the simulation running?")
        return
    
    # Initialize ROS 2
    rclpy.init()
    
    viewer = None
    try:
        viewer = CameraViewer(
            topic=args.topic,
            fps=args.fps,
            window_name=args.window_name,
            record_file=args.record,
            show_info=not args.no_info
        )
        
        rclpy.spin(viewer)
        
    except KeyboardInterrupt:
        print('\nShutting down...')
    finally:
        if viewer:
            viewer.shutdown()
            viewer.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()
        cv2.destroyAllWindows()


if __name__ == '__main__':
    main()