"""
PiPedal WebSocket proxy.

Connects to PiPedal's WebSocket at ws://<host>:8080/pipedal and relays
messages to/from the browser via FastAPI's WebSocket endpoint.

Protocol (PiPedal WS):
    Messages are JSON arrays: [{"message": "msgName", "replyTo": N}, optionalBody]
    Replies: [{"reply": N, "message": "msgName"}, optionalBody]
"""

import asyncio
import json
import logging
from typing import Optional

import websockets

logger = logging.getLogger(__name__)


class PiPedalWSProxy:
    """Manages a single upstream WebSocket connection to PiPedal and relays
    messages to a single downstream (browser) WebSocket."""

    def __init__(self, host: str = "192.168.0.245", port: int = 8080, path: str = "/pipedal"):
        self._host = host
        self._port = port
        self._path = path
        self._upstream_ws: Optional[websockets.WebSocketClientProtocol] = None
        self._downstream_ws: Optional["WebSocket"] = None  # fastapi WebSocket
        self._task: Optional[asyncio.Task] = None
        self._connected = False
        self._lock = asyncio.Lock()

    @property
    def upstream_url(self) -> str:
        return f"ws://{self._host}:{self._port}{self._path}"

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

    async def connect_upstream(self) -> bool:
        """Connect to PiPedal's WebSocket. Returns True on success."""
        try:
            self._upstream_ws = await websockets.connect(
                self.upstream_url,
                ping_interval=30,
                ping_timeout=10,
                close_timeout=5,
            )
            self._connected = True
            logger.info("Connected to PiPedal WS at %s", self.upstream_url)
            return True
        except Exception as e:
            logger.warning("Failed to connect to PiPedal WS: %s", e)
            self._connected = False
            return False

    async def close_upstream(self):
        """Close upstream connection."""
        self._connected = False
        if self._upstream_ws:
            try:
                await self._upstream_ws.close()
            except Exception:
                pass
            self._upstream_ws = None

    async def relay_loop(self, downstream_ws):
        """Run relay loop: forward upstream→downstream and downstream→upstream."""
        self._downstream_ws = downstream_ws

        async def upstream_to_downstream():
            """Read from PiPedal WS and forward to browser."""
            try:
                async for message in self._upstream_ws:
                    try:
                        await self._downstream_ws.send_text(message)
                    except Exception:
                        logger.warning("Downstream send failed, stopping relay")
                        break
            except websockets.exceptions.ConnectionClosed:
                logger.info("Upstream connection closed")
            except Exception as e:
                logger.warning("Upstream→downstream relay error: %s", e)
            finally:
                self._connected = False

        async def downstream_to_upstream():
            """Read from browser and forward to PiPedal WS."""
            try:
                while True:
                    raw = await self._downstream_ws.receive_text()
                    if self._upstream_ws and self._connected:
                        await self._upstream_ws.send(raw)
            except Exception:
                pass
            finally:
                self._connected = False

        # Run both directions concurrently
        await asyncio.gather(
            upstream_to_downstream(),
            downstream_to_upstream(),
        )

    async def send_to_pipedal(self, message: str):
        """Send a raw message to PiPedal upstream."""
        async with self._lock:
            if self._upstream_ws and self._connected:
                await self._upstream_ws.send(message)

    async def request(self, message_name: str, body=None) -> dict:
        """Send a request to PiPedal and wait for reply using replyTo/reply mechanism."""
        reply_to = id(message_name) & 0xFFFFFFFF  # simple unique id
        msg = [{"message": message_name, "replyTo": reply_to}]
        if body is not None:
            msg.append(body)

        future: asyncio.Future = asyncio.get_running_loop().create_future()
        # We need a reply handler - but since this is a simple proxy,
        # we use the REST endpoint approach instead for discovery.
        # For inline use, we'd need pending_request tracking.
        await self.send_to_pipedal(json.dumps(msg))
        # Note: real reply matching requires tracking in the relay loop.
        # This is a simplified version.
        return {"status": "sent", "reply_to": reply_to}