"""Stream input — receive JPEG frames from a remote DJ HUD via WebSocket.

This is the inverse of vdoninja.py: instead of *pushing* frames to browsers,
we *pull* frames from another DJ HUD instance's /ws endpoint.

    Remote DJ HUD (10.66.66.3:7088)
        │  binary JPEG frames over WebSocket
        ▼
    StreamReceiver (this module)
        │  latest_frame: PIL.Image (non-blocking read)
        ▼
    compose_canvas() crops regions from it

Each StreamReceiver runs an asyncio event loop in a daemon thread,
same pattern as VDONinjaServer.
"""

import asyncio
import io
import threading
import time
from typing import Optional

try:
    import aiohttp
except ImportError:
    aiohttp = None

try:
    from PIL import Image
except ImportError:
    Image = None


class StreamReceiver:
    """Connect to a remote DJ HUD WebSocket and hold the latest frame.

    Usage:
        rx = StreamReceiver("ws://10.66.66.3:7088/ws")
        rx.start()
        ...
        frame = rx.frame  # PIL.Image or None (non-blocking)
        ...
        rx.stop()
    """

    def __init__(self, ws_url: str, reconnect_interval: float = 2.0):
        self.ws_url = ws_url
        self.reconnect_interval = reconnect_interval
        self._latest_frame: Optional["Image.Image"] = None
        self._lock = threading.Lock()
        self._running = False
        self._connected = False
        self._thread: Optional[threading.Thread] = None
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._last_frame_time: float = 0.0
        self._frame_count: int = 0
        self._fps: float = 0.0
        self._fps_window_start: float = 0.0
        self._fps_window_count: int = 0

    # ── Public API ───────────────────────────────────────────────────

    def start(self):
        """Start receiving frames in a background thread."""
        if self._running:
            return
        if aiohttp is None:
            raise ImportError(
                "aiohttp is required for stream input. "
                "Install it with: pip install aiohttp"
            )
        self._running = True
        self._thread = threading.Thread(
            target=self._run_loop, daemon=True, name=f"StreamRx-{self.ws_url}"
        )
        self._thread.start()

    def stop(self):
        """Stop the receiver and clean up."""
        self._running = False
        if self._loop and not self._loop.is_closed():
            self._loop.call_soon_threadsafe(self._loop.stop)
        if self._thread:
            self._thread.join(timeout=3.0)
            self._thread = None
        self._connected = False

    @property
    def frame(self) -> Optional["Image.Image"]:
        """Get the most recently received frame (non-blocking)."""
        with self._lock:
            return self._latest_frame

    @property
    def connected(self) -> bool:
        return self._connected

    @property
    def fps(self) -> float:
        return self._fps

    @property
    def frame_count(self) -> int:
        return self._frame_count

    # ── Internal ─────────────────────────────────────────────────────

    def _run_loop(self):
        """Run asyncio event loop in dedicated thread."""
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)
        try:
            self._loop.run_until_complete(self._receive_loop())
        except Exception:
            pass
        finally:
            try:
                self._loop.close()
            except Exception:
                pass
            self._loop = None

    async def _receive_loop(self):
        """Connect, receive frames, auto-reconnect on failure."""
        while self._running:
            try:
                await self._connect_and_receive()
            except Exception as e:
                print(f"[StreamRx] Connection error ({self.ws_url}): {e}")
            self._connected = False
            if self._running:
                await asyncio.sleep(self.reconnect_interval)

    async def _connect_and_receive(self):
        """Single connection session — receive until disconnect."""
        timeout = aiohttp.ClientTimeout(total=None, sock_connect=5.0)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            print(f"[StreamRx] Connecting to {self.ws_url}...")
            async with session.ws_connect(self.ws_url) as ws:
                self._connected = True
                self._fps_window_start = time.monotonic()
                self._fps_window_count = 0
                print(f"[StreamRx] Connected to {self.ws_url}")

                async for msg in ws:
                    if not self._running:
                        break
                    if msg.type == aiohttp.WSMsgType.BINARY:
                        self._decode_frame(msg.data)
                    elif msg.type in (
                        aiohttp.WSMsgType.ERROR,
                        aiohttp.WSMsgType.CLOSE,
                        aiohttp.WSMsgType.CLOSING,
                        aiohttp.WSMsgType.CLOSED,
                    ):
                        break

                print(f"[StreamRx] Disconnected from {self.ws_url}")

    def _decode_frame(self, data: bytes):
        """Decode JPEG bytes to PIL Image and store as latest frame."""
        if Image is None:
            return
        try:
            img = Image.open(io.BytesIO(data))
            img.load()  # force decode now
            with self._lock:
                self._latest_frame = img.convert("RGB")
            self._frame_count += 1
            self._last_frame_time = time.monotonic()

            # FPS calculation (rolling 1-second window)
            self._fps_window_count += 1
            elapsed = self._last_frame_time - self._fps_window_start
            if elapsed >= 1.0:
                self._fps = self._fps_window_count / elapsed
                self._fps_window_start = self._last_frame_time
                self._fps_window_count = 0
        except Exception:
            pass  # corrupted frame, skip it
