stockautomtion/app/services/realtime.py

from __future__ import annotations

import asyncio
import json
from typing import Callable

import websockets

from app.core.auth import token_manager
from app.core.config import get_ws_url, settings
from app.core.logger import setup_logger
from app.core.rate_limiter import RateLimiter

logger = setup_logger("realtime")

_SUBSCRIBE_TR_IDS = {
    "price": "H0STCNT0",
    "orderbook": "H0STASP0",
    "trade": "H0STCNI0",
}


class RealtimeService:
    def __init__(self) -> None:
        self._ws: websockets.WebSocketClientProtocol | None = None
        self._running = False
        self._callbacks: dict[str, list[Callable]] = {}
        self._subscriptions: dict[str, list[str]] = {}
        self._rate_limiter = RateLimiter(requests_per_second=0.5)
        self._heartbeat_task: asyncio.Task | None = None
        self._receive_task: asyncio.Task | None = None

    @property
    def is_connected(self) -> bool:
        return self._ws is not None and self._running

    def on(self, event: str, callback: Callable) -> None:
        self._callbacks.setdefault(event, []).append(callback)

    async def connect(self) -> None:
        approval_key = await token_manager.get_approval_key()
        url = get_ws_url()

        try:
            self._ws = await websockets.connect(
                url,
                additional_headers={"approval_key": approval_key, "type": "Y"},
            )
            self._running = True
            self._receive_task = asyncio.create_task(self._receive_loop())
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
            logger.info("WebSocket 연결 성공: %s", url)
        except Exception as e:
            logger.error("WebSocket 연결 실패: %s", e)
            raise

    async def subscribe(self, stock_code: str, data_type: str = "price") -> None:
        if not self._ws:
            await self.connect()

        tr_id = _SUBSCRIBE_TR_IDS.get(data_type, "H0STCNT0")
        msg = {
            "header": {
                "approval_key": await token_manager.get_approval_key(),
                "custtype": "P",
                "tr_type": "1",
                "content-type": "utf-8",
            },
            "body": {
                "input": {
                    "tr_id": tr_id,
                    "tr_key": stock_code,
                }
            },
        }

        await self._ws.send(json.dumps(msg))
        self._subscriptions.setdefault(data_type, []).append(stock_code)
        logger.info("구독 요청: %s - %s", data_type, stock_code)

    async def unsubscribe(self, stock_code: str, data_type: str = "price") -> None:
        if not self._ws:
            return

        tr_id = _SUBSCRIBE_TR_IDS.get(data_type, "H0STCNT0")
        msg = {
            "header": {
                "approval_key": await token_manager.get_approval_key(),
                "custtype": "P",
                "tr_type": "2",
                "content-type": "utf-8",
            },
            "body": {
                "input": {
                    "tr_id": tr_id,
                    "tr_key": stock_code,
                }
            },
        }

        await self._ws.send(json.dumps(msg))
        subs = self._subscriptions.get(data_type, [])
        if stock_code in subs:
            subs.remove(stock_code)
        logger.info("구독 해제: %s - %s", data_type, stock_code)

    async def _receive_loop(self) -> None:
        try:
            async for raw_msg in self._ws:
                await self._handle_message(raw_msg)
        except websockets.ConnectionClosed:
            logger.warning("WebSocket 연결 끊김. 5초 후 재연결...")
        except Exception as e:
            logger.error("WebSocket 수신 오류: %s", e)
        if self._running:
            await asyncio.sleep(5)
            try:
                await self.connect()
            except Exception:
                logger.error("WebSocket 재연결 실패")

    async def _handle_message(self, raw_msg: str) -> None:
        try:
            msg = json.loads(raw_msg)
            header = msg.get("header", {})
            body = msg.get("body", {})

            tr_id = header.get("tr_id", "")

            if tr_id == "PINGPONG":
                return

            parsed = self._parse_data(tr_id, body.get("output", {}))
            if not parsed:
                return

            event_type = "price"
            if "bid_prices" in parsed:
                event_type = "orderbook"
            elif "trade_price" in parsed:
                event_type = "trade"

            for callback in self._callbacks.get(event_type, []):
                try:
                    result = callback(parsed)
                    if asyncio.iscoroutine(result):
                        await result
                except Exception as e:
                    logger.error("콜백 실행 오류: %s", e)

        except json.JSONDecodeError:
            logger.warning("JSON 파싱 실패")

    def _parse_data(self, tr_id: str, output: dict) -> dict | None:
        if not output:
            return None

        if tr_id == "H0STCNT0":
            return {
                "type": "price",
                "stock_code": output.get("mksc_shrn_iscd", ""),
                "trade_price": int(output.get("stck_prpr", 0)),
                "change_price": int(output.get("prdy_vrss", 0)),
                "change_rate": float(output.get("prdy_ctrt", 0)),
                "open_price": int(output.get("stck_oprc", 0)),
                "high_price": int(output.get("stck_hgpr", 0)),
                "low_price": int(output.get("stck_lwpr", 0)),
                "volume": int(output.get("acml_vol", 0)),
                "trade_time": output.get("stck_cntg_hour", ""),
            }

        if tr_id == "H0STASP0":
            return {
                "type": "orderbook",
                "stock_code": output.get("mksc_shrn_iscd", ""),
                "bid_prices": [int(output.get(f"phsc_kprc_{i}", 0)) for i in range(1, 6)],
                "ask_prices": [int(output.get(f"sats_kprc_{i}", 0)) for i in range(1, 6)],
                "bid_volumes": [int(output.get(f"phsc_vola_{i}", 0)) for i in range(1, 6)],
                "ask_volumes": [int(output.get(f"sats_ac_vola_{i}", 0)) for i in range(1, 6)],
            }

        return None

    async def _heartbeat_loop(self) -> None:
        while self._running:
            try:
                await asyncio.sleep(30)
                if self._ws and self._ws.open:
                    await self._ws.send(json.dumps({"header": {"tr_id": "PINGPONG"}}))
            except Exception:
                break

    async def disconnect(self) -> None:
        self._running = False
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
        if self._receive_task:
            self._receive_task.cancel()
        if self._ws:
            await self._ws.close()
            self._ws = None
        logger.info("WebSocket 연결 종료")


realtime_service = RealtimeService()