"""Stream client — 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
        ▼
    StreamClient (this module)
        │  get_frame() → PIL.Image (non-blocking, thread-safe)
        ▼
    compose_canvas() crops regions from it

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

Supported URL formats:
    ws://10.66.66.3:7088/ws          — direct WebSocket URL
    http://10.66.66.3:7088/viewer    — HTTP viewer URL (auto-converts to ws://)
    http://10.66.66.3:7088           — server root (auto-appends /ws)
    10.66.66.3:7088                  — bare host:port (auto-prefixes ws://)
"""

import asyncio
import io
import threading
import time
from typing import Optional
from urllib.parse import urlparse

try:
    import aiohttp
except ImportError:
    aiohttp = None

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


def normalize_stream_url(url: str) -> str:
    """Convert various URL formats to a ws:// WebSocket URL.

    Examples:
        ws://10.66.66.3:7088/ws          → ws://10.66.66.3:7088/ws
        wss://host/ws                     → wss://host/ws
        http://10.66.66.3:7088/viewer     → ws://10.66.66.3:7088/ws
        http://10.66.66.3:7088            → ws://10.66.66.3:7088/ws
        https://host:7088/viewer          → wss://host:7088/ws
        10.66.66.3:7088                   → ws://10.66.66.3:7088/ws
    """
    url = url.strip()
    if not url:
        return url

    # Bare host:port (no scheme)
    if "://" not in url:
        url = "ws://" + url

    parsed = urlparse(url)

    # Convert http(s) → ws(s)
    scheme = parsed.scheme
    if scheme == "http":
        scheme = "ws"
    elif scheme == "https":
        scheme = "wss"

    # Ensure path ends at /ws
    path = parsed.path.rstrip("/")
    if path != "/ws":
        path = "/ws"

    return f"{scheme}://{parsed.netloc}{path}"


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

    Usage:
        client = StreamClient("ws://10.66.66.3:7088/ws")
        client.start()
        ...
        frame = client.get_frame()  # PIL Image or None
        ...
        client.stop()

    Thread-safe: start/stop/get_frame can be called from any thread.
    The WebSocket connection runs in a dedicated daemon thread.
    """

    def __init__(self, url: str, reconnect_delay: float = 2.0):
        if aiohttp is None:
            raise ImportError(
                "aiohttp is required for stream input. "
                "Install it with: pip install aiohttp"
            )
        self._raw_url = url
        self._ws_url = normalize_stream_url(url)
        self._reconnect_delay = reconnect_delay

        self._frame: Optional["Image.Image"] = None
        self._frame_lock = threading.Lock()
        self._frame_size: Optional[tuple] = None  # (width, height) of last frame

        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._error: str = ""

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

    def start(self):
        """Start the background WebSocket receiver thread."""
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(
            target=self._run_loop, daemon=True,
            name=f"StreamClient-{self._ws_url}"
        )
        self._thread.start()

    def stop(self):
        """Stop the receiver and close the WebSocket."""
        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

    def get_frame(self) -> Optional["Image.Image"]:
        """Return the latest frame as a PIL Image, or None if no frame yet."""
        with self._frame_lock:
            return self._frame

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

    @property
    def is_running(self) -> bool:
        return self._running

    @property
    def url(self) -> str:
        return self._ws_url

    @property
    def raw_url(self) -> str:
        return self._raw_url

    @property
    def frame_size(self) -> Optional[tuple]:
        """(width, height) of the last received frame, or None."""
        return self._frame_size

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

    @property
    def error(self) -> str:
        return self._error

    # ── Background thread ────────────────────────────────────────────

    def _run_loop(self):
        """Run the asyncio event loop in the background thread."""
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)
        try:
            self._loop.run_until_complete(self._receive_loop())
        except Exception as e:
            self._error = str(e)
            print(f"[StreamClient] Fatal error: {e}")
        finally:
            try:
                self._loop.close()
            except Exception:
                pass
            self._loop = None
            self._connected = False

    async def _receive_loop(self):
        """Connect, receive frames, reconnect on failure."""
        while self._running:
            try:
                await self._connect_and_receive()
            except asyncio.CancelledError:
                break
            except Exception as e:
                self._error = str(e)
                self._connected = False
                if self._running:
                    print(f"[StreamClient] Connection lost ({self._ws_url}): {e}")
                    print(f"[StreamClient] Reconnecting in {self._reconnect_delay}s...")
                    await asyncio.sleep(self._reconnect_delay)

    async def _connect_and_receive(self):
        """Single connection attempt: connect, receive until disconnected."""
        self._error = ""
        timeout = aiohttp.ClientTimeout(total=None, connect=10)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            print(f"[StreamClient] Connecting to {self._ws_url}...")
            async with session.ws_connect(
                self._ws_url,
                heartbeat=20.0,
                autoping=True,
            ) as ws:
                self._connected = True
                print(f"[StreamClient] Connected to {self._ws_url}")
                self._error = ""

                async for msg in ws:
                    if not self._running:
                        break
                    if msg.type == aiohttp.WSMsgType.BINARY:
                        self._decode_frame(msg.data)
                    elif msg.type == aiohttp.WSMsgType.ERROR:
                        self._error = f"WebSocket error: {ws.exception()}"
                        break
                    elif msg.type in (
                        aiohttp.WSMsgType.CLOSE,
                        aiohttp.WSMsgType.CLOSING,
                        aiohttp.WSMsgType.CLOSED,
                    ):
                        break

        self._connected = False

    def _decode_frame(self, data: bytes):
        """Decode JPEG bytes into a PIL Image and store as latest frame."""
        if Image is None:
            return
        try:
            img = Image.open(io.BytesIO(data))
            img.load()  # force decode
            with self._frame_lock:
                self._frame = img
                self._frame_size = img.size
            self._frame_count += 1
            self._last_frame_time = time.monotonic()
        except Exception as e:
            # Corrupted frame — skip it, keep the previous frame
            pass

    # ── Representation ───────────────────────────────────────────────

    def __repr__(self):
        state = "connected" if self._connected else "disconnected"
        size = f"{self._frame_size[0]}×{self._frame_size[1]}" if self._frame_size else "?"
        return f"StreamClient({self._ws_url!r}, {state}, {size}, {self._frame_count} frames)"
