# Copyright (c) 2025, Yutang Li, SIAT (yt.li2@siat.ac.cn)
# SPDX-License-Identifier: BSD-3-Clause

"""Sim-to-VR H264 stereo streaming manager."""

import logging

logger = logging.getLogger(__name__)


class StreamingManager:
    """Manages stereo camera → H264 → TCP streaming to VR headset.

    Wraps SimVideoStreamer with graceful error handling and auto-disable on failure.
    """

    def __init__(self, host: str, port: int, scene, bitrate: int = 20_000_000):
        self._streamer = None
        self._enabled = False

        if "vr_left_eye" not in scene.sensors:
            logger.warning("[Stream] vr_left_eye camera not found. Streaming disabled.")
            return

        try:
            from mindbot.utils.sim_video_streamer import SimVideoStreamer
            cam = scene["vr_left_eye"]
            self._streamer = SimVideoStreamer(
                host=host, port=port,
                width=cam.cfg.width, height=cam.cfg.height,
                fps=30, bitrate=bitrate,
            )
            if self._streamer.connect():
                self._enabled = True
                print(f"[INFO] Sim stereo streaming to {host}:{port} "
                      f"({cam.cfg.width * 2}x{cam.cfg.height} SBS @ 30fps)")
            else:
                logger.error("[Stream] Failed to connect. Streaming disabled.")
                self._streamer = None
        except ImportError:
            logger.error("[Stream] mindbot.utils.sim_video_streamer not found. pip install av?")
        except Exception as e:
            logger.error(f"[Stream] Init failed: {e}")

    @property
    def enabled(self) -> bool:
        return self._enabled

    _debug_saved = False

    def send(self, scene, frame_count: int = 0) -> None:
        """Send one stereo frame from scene cameras. Auto-disables on failure."""
        if not self._enabled:
            return
        try:
            left_rgb = scene["vr_left_eye"].data.output["rgb"][0].cpu().numpy()
            right_rgb = scene["vr_right_eye"].data.output["rgb"][0].cpu().numpy()

            # Save debug images once (frame 5 to ensure rendering is stable)
            if frame_count == 150 and not self._debug_saved:
                self._debug_saved = True
                try:
                    from PIL import Image
                    import numpy as np
                    Image.fromarray(left_rgb[..., :3]).save("/tmp/vr_left_debug.png")
                    Image.fromarray(right_rgb[..., :3]).save("/tmp/vr_right_debug.png")
                    sbs = np.concatenate([left_rgb, right_rgb], axis=1)
                    Image.fromarray(sbs[..., :3]).save("/tmp/vr_sbs_debug.png")
                    print(f"[STREAM DIAG] Saved debug frames: left={left_rgb.shape}, right={right_rgb.shape}")
                except Exception as e:
                    print(f"[STREAM DIAG] Save failed: {e}")

            if not self._streamer.send_frame(left_rgb, right_rgb):
                logger.warning("[Stream] Connection lost. Disabling streaming.")
                self._enabled = False
        except Exception as e:
            if frame_count % 300 == 0:
                logger.warning(f"[Stream] Frame error: {e}")

    def close(self):
        if self._streamer is not None:
            self._streamer.close()