"""
WebSocket Push Service — real-time streaming for mobile & desktop.

Replaces 12 polling timers with a single WebSocket connection.
Broadcasts: price ticks, position changes, risk alerts, signal updates.

Usage in gui_backend.py:
    from src.ws_push import ws_router, ws_manager
    app.include_router(ws_router)
    # Then call ws_manager.broadcast("price", data) from any handler
"""

from __future__ import annotations

import asyncio
import json
import time
from typing import Any, Dict, List, Optional, Set

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

ws_router = APIRouter(prefix="/ws", tags=["websocket"])


class WebSocketManager:
    """Manages WebSocket connections and broadcast."""

    def __init__(self):
        self._connections: Set[WebSocket] = set()
        self._lock = asyncio.Lock()

    async def connect(self, ws: WebSocket) -> None:
        await ws.accept()
        async with self._lock:
            self._connections.add(ws)

    async def disconnect(self, ws: WebSocket) -> None:
        async with self._lock:
            self._connections.discard(ws)

    async def broadcast(self, event_type: str, data: Dict[str, Any]) -> None:
        """Push an event to all connected clients."""
        message = json.dumps({
            "type": event_type,
            "data": data,
            "ts": time.time(),
        }, ensure_ascii=False, default=str)

        dead: List[WebSocket] = []
        async with self._lock:
            connections = list(self._connections)

        for ws in connections:
            try:
                await ws.send_text(message)
            except Exception:
                dead.append(ws)

        if dead:
            async with self._lock:
                for ws in dead:
                    self._connections.discard(ws)

    async def broadcast_price(self, bid: float, ask: float, spread: float, symbol: str = "XAUUSD") -> None:
        """Push a price tick."""
        await self.broadcast("price", {
            "symbol": symbol,
            "bid": bid,
            "ask": ask,
            "spread": spread,
        })

    async def broadcast_position(self, positions: List[Dict[str, Any]]) -> None:
        """Push position updates."""
        await self.broadcast("positions", {
            "count": len(positions),
            "positions": [
                {
                    "ticket": p.get("ticket"),
                    "symbol": p.get("symbol", "XAUUSD"),
                    "side": p.get("side"),
                    "volume": p.get("volume"),
                    "entry": p.get("entry"),
                    "pnl": p.get("pnl"),
                }
                for p in positions
            ],
        })

    async def broadcast_risk(self, light: str, risk_code: str, decision: str) -> None:
        """Push risk/status light changes."""
        await self.broadcast("risk", {
            "light": light,
            "risk_code": risk_code,
            "decision": decision,
        })

    async def broadcast_signal(self, signal: Dict[str, Any]) -> None:
        """Push new trading signal."""
        await self.broadcast("signal", {
            "action": signal.get("action"),
            "side": signal.get("side"),
            "strategy": signal.get("strategy_id"),
            "confidence": signal.get("confidence"),
        })

    async def broadcast_alert(self, level: str, message: str, category: str = "system") -> None:
        """Push an alert/notification."""
        await self.broadcast("alert", {
            "level": level,  # info / warning / critical
            "message": message,
            "category": category,
        })

    @property
    def connection_count(self) -> int:
        return len(self._connections)


# ── Singleton ──

ws_manager = WebSocketManager()


# ── WebSocket endpoint ──

@ws_router.websocket("/live")
async def ws_live_endpoint(websocket: WebSocket):
    """Main WebSocket endpoint for real-time data.

    Client can send JSON messages to subscribe to specific channels:
      {"subscribe": ["price", "positions", "risk", "signals", "alerts"]}

    Server pushes events of subscribed types automatically.
    """
    await ws_manager.connect(websocket)
    subscriptions: Set[str] = {"price", "positions", "risk", "alerts"}  # default: all

    try:
        # Send welcome + connection count
        await websocket.send_text(json.dumps({
            "type": "connected",
            "data": {"message": "WebSocket connected", "clients": ws_manager.connection_count},
            "ts": time.time(),
        }))

        while True:
            try:
                raw = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
                msg = json.loads(raw)
                if isinstance(msg.get("subscribe"), list):
                    subscriptions = set(msg["subscribe"])
                elif msg.get("type") == "ping":
                    await websocket.send_text(json.dumps({"type": "pong", "ts": time.time()}))
            except asyncio.TimeoutError:
                # Send keepalive ping
                try:
                    await websocket.send_text(json.dumps({"type": "ping", "ts": time.time()}))
                except Exception:
                    break
            except (json.JSONDecodeError, KeyError):
                pass

    except (WebSocketDisconnect, Exception):
        pass
    finally:
        await ws_manager.disconnect(websocket)


@ws_router.get("/status")
async def ws_status():
    """Get WebSocket server status."""
    return {
        "connections": ws_manager.connection_count,
        "endpoint": "/ws/live",
    }
