from __future__ import annotations

import asyncio
import json

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.core.logger import setup_logger
from app.services.realtime import realtime_service

router = APIRouter(tags=["websocket"])
logger = setup_logger("ws_router")

_connected_clients: set[WebSocket] = set()


async def _broadcast(data: dict) -> None:
    message = json.dumps(data, ensure_ascii=False)
    disconnected: list[WebSocket] = []
    for client in _connected_clients:
        try:
            await client.send_text(message)
        except Exception:
            disconnected.append(client)
    for client in disconnected:
        _connected_clients.discard(client)


@router.websocket("/ws/realtime")
async def websocket_endpoint(websocket: WebSocket) -> None:
    await websocket.accept()
    _connected_clients.add(websocket)
    logger.info("WebSocket 클라이언트 연결 (총 %d)", len(_connected_clients))

    realtime_service.on("price", _broadcast)
    realtime_service.on("orderbook", _broadcast)

    try:
        while True:
            raw = await websocket.receive_text()
            try:
                msg = json.loads(raw)
                action = msg.get("action", "")

                if action == "subscribe":
                    stock_code = msg.get("stock_code", "")
                    data_type = msg.get("data_type", "price")
                    await realtime_service.subscribe(stock_code, data_type)
                    await websocket.send_text(
                        json.dumps({"status": "subscribed", "stock_code": stock_code, "data_type": data_type})
                    )

                elif action == "unsubscribe":
                    stock_code = msg.get("stock_code", "")
                    data_type = msg.get("data_type", "price")
                    await realtime_service.unsubscribe(stock_code, data_type)
                    await websocket.send_text(
                        json.dumps({"status": "unsubscribed", "stock_code": stock_code})
                    )

            except json.JSONDecodeError:
                pass

    except WebSocketDisconnect:
        pass
    finally:
        _connected_clients.discard(websocket)
        logger.info("WebSocket 클라이언트 해제 (총 %d)", len(_connected_clients))