# SPDX-FileCopyrightText: Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# SPDX-License-Identifier: BSD-3-Clause

"""Simulation stereo video streamer for VR headset.

Encodes stereo RGB frames (left + right eye) as H264 side-by-side video
and streams to the Unity VR client over TCP, using the same 4-byte
big-endian length-prefixed NALU protocol as RobotVisionConsole.

Usage:
    streamer = SimVideoStreamer("172.20.103.171", 12345, width=960, height=540)
    streamer.connect()
    # In the sim loop:
    streamer.send_frame(left_rgb_np, right_rgb_np)
    # Cleanup:
    streamer.close()

The Unity client (PICO) acts as TCP SERVER on the specified port.
This streamer connects as TCP CLIENT — same role as RobotVisionConsole.
"""

from __future__ import annotations

import logging
import socket
import struct
import threading
import time
from typing import Optional

import av
import numpy as np

logger = logging.getLogger(__name__)


class SimVideoStreamer:
    """H264 stereo video streamer over TCP for VR headset display.

    Accepts left/right RGB numpy arrays, assembles a side-by-side frame,
    encodes to H264 with PyAV, and sends each NALU packet prefixed with
    a 4-byte big-endian length header over a TCP socket.
    """

    def __init__(
        self,
        host: str,
        port: int = 12345,
        width: int = 960,
        height: int = 540,
        fps: int = 30,
        bitrate: int = 20_000_000,
    ):
        """Initialize the streamer.

        Args:
            host: IP address of the PICO VR headset (TCP server).
            port: TCP port the PICO is listening on.
            width: Width of EACH eye image (SBS output = width*2 x height).
            height: Height of each eye image.
            fps: Target frame rate for H264 encoding.
            bitrate: H264 encoding bitrate in bits per second.
        """
        self.host = host
        self.port = port
        self.eye_width = width
        self.eye_height = height
        self.sbs_width = width * 2
        self.sbs_height = height
        self.fps = fps
        self.bitrate = bitrate

        self._sock: Optional[socket.socket] = None
        self._codec_ctx: Optional[av.CodecContext] = None
        self._frame_idx = 0
        self._last_frame_time = 0.0
        self._min_frame_interval = 1.0 / fps

        # Async encoding: background thread picks up frames from _pending_frame
        self._pending_frame: Optional[np.ndarray] = None
        self._frame_lock = threading.Lock()
        self._encode_thread: Optional[threading.Thread] = None
        self._running = False

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def connect(self) -> bool:
        """Connect to the PICO VR headset TCP server.

        Returns:
            True if connection succeeded.
        """
        try:
            self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            self._sock.connect((self.host, self.port))
            logger.info(f"[SimVideoStreamer] Connected to {self.host}:{self.port}")
        except OSError as e:
            logger.error(f"[SimVideoStreamer] TCP connect failed: {e}")
            self._sock = None
            return False

        self._init_encoder()
        self._frame_idx = 0
        self._last_frame_time = 0.0

        # Start background encode+send thread
        self._running = True
        self._encode_thread = threading.Thread(target=self._encode_loop, daemon=True)
        self._encode_thread.start()
        return True

    def send_frame(self, left_rgb: np.ndarray, right_rgb: np.ndarray) -> bool:
        """Submit a stereo frame for async encoding and sending.

        This method is NON-BLOCKING: it copies the frame data and returns
        immediately. The background thread handles encoding and TCP send.
        If the previous frame hasn't been encoded yet, this frame replaces it
        (latest-frame-wins, no queue buildup).

        Args:
            left_rgb: Left eye image, shape (H, W, 3), dtype uint8, RGB order.
            right_rgb: Right eye image, shape (H, W, 3), dtype uint8, RGB order.

        Returns:
            True if frame was accepted. False if connection is lost.
        """
        if not self._running:
            return False

        # Assemble SBS and hand off to background thread (latest-frame-wins)
        sbs = np.concatenate([left_rgb, right_rgb], axis=1)  # (H, W*2, 3)
        with self._frame_lock:
            self._pending_frame = sbs

        return True

    def close(self):
        """Stop background thread, flush encoder, and close connection."""
        self._running = False
        if self._encode_thread is not None:
            self._encode_thread.join(timeout=2.0)
            self._encode_thread = None

        if self._codec_ctx is not None:
            try:
                packets = self._codec_ctx.encode(None)
                if self._sock is not None:
                    for pkt in packets:
                        nalu_data = bytes(pkt)
                        header = struct.pack(">I", len(nalu_data))
                        self._sock.sendall(header + nalu_data)
            except Exception:
                pass
            self._codec_ctx = None

        if self._sock is not None:
            try:
                self._sock.close()
            except Exception:
                pass
            self._sock = None
            logger.info("[SimVideoStreamer] Connection closed.")

    # ------------------------------------------------------------------
    # Private
    # ------------------------------------------------------------------

    def _init_encoder(self):
        """Initialize H264 encoder via PyAV."""
        codec = av.Codec("libx264", "w")
        self._codec_ctx = codec.create()
        self._codec_ctx.width = self.sbs_width
        self._codec_ctx.height = self.sbs_height
        self._codec_ctx.pix_fmt = "yuv420p"
        from fractions import Fraction
        self._codec_ctx.time_base = Fraction(1, self.fps)
        self._codec_ctx.framerate = Fraction(self.fps, 1)
        self._codec_ctx.bit_rate = self.bitrate
        self._codec_ctx.gop_size = self.fps  # keyframe every 1 second
        self._codec_ctx.max_b_frames = 0  # no B-frames for low latency
        self._codec_ctx.thread_count = 2

        # Low-latency encoding options
        self._codec_ctx.options = {
            "preset": "ultrafast",
            "tune": "zerolatency",
            "profile": "baseline",  # widest decoder compatibility
        }

        self._codec_ctx.open()
        logger.info(
            f"[SimVideoStreamer] H264 encoder initialized: "
            f"{self.sbs_width}x{self.sbs_height} @ {self.fps}fps, "
            f"bitrate={self.bitrate/1e6:.1f}Mbps"
        )

    def _encode_loop(self):
        """Background thread: encode and send frames at target fps."""
        while self._running:
            # Grab the latest pending frame (if any)
            with self._frame_lock:
                sbs = self._pending_frame
                self._pending_frame = None

            if sbs is None:
                time.sleep(0.001)  # 1ms idle wait
                continue

            # Rate limit encoding to target fps
            now = time.monotonic()
            if now - self._last_frame_time < self._min_frame_interval:
                continue
            self._last_frame_time = now

            # Encode
            frame = av.VideoFrame.from_ndarray(sbs, format="rgb24")
            frame.pts = self._frame_idx
            self._frame_idx += 1

            try:
                packets = self._codec_ctx.encode(frame)
            except av.error.FFmpegError as e:
                logger.warning(f"[SimVideoStreamer] Encode error: {e}")
                continue

            # Send each packet
            for pkt in packets:
                nalu_data = bytes(pkt)
                header = struct.pack(">I", len(nalu_data))
                try:
                    self._sock.sendall(header + nalu_data)
                except (BrokenPipeError, ConnectionResetError, OSError) as e:
                    logger.error(f"[SimVideoStreamer] Send failed: {e}")
                    self._running = False
                    return
