"""
equirect_stitcher_node.py — 360° equirectangular stitching ROS2 node.

Subscribes to 4 IMX219 camera topics, stitches to equirectangular panorama,
publishes on /camera/panoramic/image_raw and /camera/panoramic/compressed.

Pipeline:
  1. On init: load calibration YAMLs -> build EquirectProjector (precomputed remaps)
  2. Each timer tick (15 Hz): take latest frames from all cameras, call projector.stitch()
  3. Publish raw + compressed
  4. Optionally: pass through GStreamer NVENC pipeline for RTSP (see rtsp_server_node.py)

Performance targets (Jetson Orin Nano Super, 67 TOPS):
  - Remap + blend: ~10-20ms with cv2.cuda if available, ~30-50ms CPU fallback
  - NVENC encode: handled by GStreamer pipeline, <5ms
  - Target: 15 FPS at 1920x960
"""

import json
import time

import numpy as np
import cv2
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from sensor_msgs.msg import Image, CompressedImage
from std_msgs.msg import String
from cv_bridge import CvBridge

from saltybot_panoramic.calibration_loader import load_all_camera_infos
from saltybot_panoramic.equirect_projector import CameraConfig, EquirectProjector

# Camera names and their yaw directions
_CAM_NAMES = ['front', 'right', 'rear', 'left']
_CAM_YAWS  = [0.0, -90.0, 180.0, 90.0]  # degrees: 0=forward, CCW positive
_CAM_TOPICS = {
    'front': '/camera/front/image_raw',
    'right': '/camera/right/image_raw',
    'rear':  '/camera/rear/image_raw',
    'left':  '/camera/left/image_raw',
}


class EquirectStitcherNode(Node):
    """360° equirectangular stitcher.

    Subscribes to 4 camera image topics, stitches to 1920x960 equirectangular
    panorama using precomputed fisheye remap maps, and publishes the result.
    Handles cameras coming online one at a time — missing cameras produce black
    regions in the output.
    """

    def __init__(self):
        super().__init__('equirect_stitcher')

        # Declare parameters
        self.declare_parameter('calibration_dir', '/mnt/nvme/saltybot/calibration')
        self.declare_parameter('output_width', 1920)
        self.declare_parameter('output_height', 960)
        self.declare_parameter('publish_rate', 15.0)
        self.declare_parameter('blend_margin_px', 30)
        self.declare_parameter('camera_height', 0.30)
        self.declare_parameter('camera_pitch_deg', 0.0)
        self.declare_parameter('publish_compressed', True)
        self.declare_parameter('jpeg_quality', 80)
        self.declare_parameter('use_cuda', True)

        # Read parameters
        self._calib_dir     = self.get_parameter('calibration_dir').value
        self._out_w         = self.get_parameter('output_width').value
        self._out_h         = self.get_parameter('output_height').value
        self._rate          = self.get_parameter('publish_rate').value
        self._blend_margin  = self.get_parameter('blend_margin_px').value
        self._pitch_deg     = self.get_parameter('camera_pitch_deg').value
        self._pub_compressed = self.get_parameter('publish_compressed').value
        self._jpeg_quality  = self.get_parameter('jpeg_quality').value
        self._use_cuda      = self.get_parameter('use_cuda').value

        self._bridge = CvBridge()
        self._latest: dict = {name: None for name in _CAM_NAMES}
        self._frame_count = 0
        self._last_log_time = time.monotonic()

        # CUDA remap maps (GpuMat), set in _try_init_cuda
        self._cuda_maps: list = []   # list of (GpuMat_x, GpuMat_y) per camera
        self._cuda_available = False

        # Build projector
        self.get_logger().info(f'Loading calibration from {self._calib_dir}')
        cam_infos = load_all_camera_infos(self._calib_dir)
        cam_configs = []
        for name, yaw in zip(_CAM_NAMES, _CAM_YAWS):
            info = cam_infos[name]
            cfg = CameraConfig(
                name=name,
                yaw_deg=yaw,
                pitch_deg=self._pitch_deg,
                K=info['K'],
                D=info['D'],
                image_size=info['image_size'],
                fov_deg=160.0,
            )
            cam_configs.append(cfg)
            status = 'OK' if info['valid'] else 'PLACEHOLDER'
            self.get_logger().info(
                f'Camera {name}: [{status}] K={info["K"][0,0]:.1f}px focal, '
                f'size={info["image_size"]}'
            )

        self._projector = EquirectProjector(
            self._out_w, self._out_h, cam_configs, self._blend_margin
        )

        # Try to init CUDA remap
        if self._use_cuda:
            self._try_init_cuda(cam_configs)

        # QoS for camera subscriptions: best-effort, keep only last frame
        best_effort_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )

        # Subscribe to all 4 camera topics
        for name, topic in _CAM_TOPICS.items():
            self.create_subscription(
                Image, topic,
                lambda msg, n=name: self._on_image(msg, n),
                best_effort_qos,
            )
            self.get_logger().info(f'Subscribed to {topic}')

        # Publishers
        self._pub_raw = self.create_publisher(
            Image, '/camera/panoramic/image_raw',
            QoSProfile(
                reliability=ReliabilityPolicy.BEST_EFFORT,
                history=HistoryPolicy.KEEP_LAST,
                depth=1,
            ),
        )

        if self._pub_compressed:
            self._pub_comp = self.create_publisher(
                CompressedImage, '/camera/panoramic/compressed',
                QoSProfile(
                    reliability=ReliabilityPolicy.BEST_EFFORT,
                    history=HistoryPolicy.KEEP_LAST,
                    depth=1,
                ),
            )
        else:
            self._pub_comp = None

        self._pub_diag = self.create_publisher(
            String, '/camera/panoramic/diagnostics',
            QoSProfile(
                reliability=ReliabilityPolicy.RELIABLE,
                history=HistoryPolicy.KEEP_LAST,
                depth=1,
            ),
        )

        # Timer
        period = 1.0 / self._rate
        self._timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f'EquirectStitcher ready: {self._out_w}x{self._out_h} @ {self._rate}Hz, '
            f'CUDA={self._cuda_available}'
        )

    def _try_init_cuda(self, cam_configs: list) -> None:
        """Attempt to upload remap maps to GPU for CUDA-accelerated remap."""
        try:
            if not hasattr(cv2, 'cuda') or not cv2.cuda.getCudaEnabledDeviceCount():
                self.get_logger().info('cv2.cuda not available — using CPU remap')
                return
            for i, _cfg in enumerate(cam_configs):
                mx, my = self._projector._maps[i]
                gpu_mx = cv2.cuda_GpuMat()
                gpu_my = cv2.cuda_GpuMat()
                gpu_mx.upload(mx)
                gpu_my.upload(my)
                self._cuda_maps.append((gpu_mx, gpu_my))
            self._cuda_available = True
            self.get_logger().info('CUDA remap maps uploaded to GPU')
        except Exception as e:
            self.get_logger().warn(f'CUDA init failed, falling back to CPU: {e}')
            self._cuda_maps = []
            self._cuda_available = False

    def _on_image(self, msg: Image, cam_name: str) -> None:
        """Store latest frame for the given camera."""
        try:
            bgr = self._bridge.imgmsg_to_cv2(msg, 'bgr8')
            self._latest[cam_name] = bgr
        except Exception as e:
            self.get_logger().warn(f'Failed to convert image from {cam_name}: {e}',
                                   throttle_duration_sec=5.0)

    def _stitch_cuda(self, images: dict) -> np.ndarray:
        """GPU-accelerated stitch using cv2.cuda.remap."""
        out = np.zeros((self._out_h, self._out_w, 3), dtype=np.float32)
        weights = self._projector._weights

        for i, cfg in enumerate(self._projector._configs):
            img = images.get(cfg.name)
            if img is None:
                continue
            try:
                gpu_src = cv2.cuda_GpuMat()
                gpu_src.upload(img.astype(np.float32))
                gpu_mx, gpu_my = self._cuda_maps[i]
                gpu_warped = cv2.cuda.remap(
                    gpu_src, gpu_mx, gpu_my,
                    interpolation=cv2.INTER_LINEAR,
                    borderMode=cv2.BORDER_CONSTANT,
                )
                warped = gpu_warped.download()
                w = weights[i][:, :, np.newaxis]
                out += warped * w
            except Exception as e:
                self.get_logger().warn(
                    f'CUDA remap failed for {cfg.name}, using CPU: {e}',
                    throttle_duration_sec=10.0,
                )
                # CPU fallback for this camera
                mx, my = self._projector._maps[i]
                warped_cpu = cv2.remap(
                    img.astype(np.float32), mx, my,
                    interpolation=cv2.INTER_LINEAR,
                    borderMode=cv2.BORDER_CONSTANT, borderValue=0,
                )
                w = weights[i][:, :, np.newaxis]
                out += warped_cpu * w

        return np.clip(out, 0, 255).astype(np.uint8)

    def _timer_callback(self) -> None:
        """Main stitching + publish callback."""
        # Count how many cameras are available
        available = [name for name, img in self._latest.items() if img is not None]

        if not available:
            # No cameras yet — skip silently
            return

        t0 = time.monotonic()

        # Stitch panorama (missing cameras yield black in their zones)
        if self._cuda_available and self._cuda_maps:
            try:
                panorama = self._stitch_cuda(self._latest)
            except Exception as e:
                self.get_logger().warn(f'CUDA stitch error, falling back: {e}')
                panorama = self._projector.stitch(self._latest)
        else:
            panorama = self._projector.stitch(self._latest)

        stitch_ms = (time.monotonic() - t0) * 1000.0

        now = self.get_clock().now().to_msg()

        # Publish raw
        img_msg = self._bridge.cv2_to_imgmsg(panorama, encoding='bgr8')
        img_msg.header.stamp = now
        img_msg.header.frame_id = 'panoramic_camera'
        self._pub_raw.publish(img_msg)

        # Publish compressed
        if self._pub_comp is not None:
            encode_params = [cv2.IMWRITE_JPEG_QUALITY, self._jpeg_quality]
            ok, jpeg_buf = cv2.imencode('.jpg', panorama, encode_params)
            if ok:
                comp_msg = CompressedImage()
                comp_msg.header.stamp = now
                comp_msg.header.frame_id = 'panoramic_camera'
                comp_msg.format = 'jpeg'
                comp_msg.data = jpeg_buf.tobytes()
                self._pub_comp.publish(comp_msg)

        self._frame_count += 1

        # Compute actual FPS over the last 30 frames
        now_mono = time.monotonic()
        elapsed = now_mono - self._last_log_time
        fps_actual = 0.0
        if elapsed > 0:
            fps_actual = 30.0 / elapsed if self._frame_count % 30 == 0 else 0.0

        # Publish diagnostics as JSON
        diag = {
            'stitch_ms': round(stitch_ms, 1),
            'cameras_available': available,
            'cameras_missing': [n for n in _CAM_NAMES if n not in available],
            'frame_count': self._frame_count,
            'cuda': self._cuda_available,
        }
        if fps_actual > 0:
            diag['fps_actual'] = round(fps_actual, 1)

        diag_msg = String()
        diag_msg.data = json.dumps(diag)
        self._pub_diag.publish(diag_msg)

        # Log FPS every 30 frames
        if self._frame_count % 30 == 0:
            fps_log = 30.0 / elapsed if elapsed > 0 else 0.0
            self.get_logger().info(
                f'Panoramic: {fps_log:.1f} FPS, stitch {stitch_ms:.0f}ms, '
                f'cams={available}, CUDA={self._cuda_available}'
            )
            self._last_log_time = now_mono


def main(args=None):
    rclpy.init(args=args)
    node = EquirectStitcherNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
