auto-stock/backend/app/realtime/websocket_client.py

"""
KIS WebSocket Client

KIS API WebSocket 연결 및 실시간 데이터 수신
"""
import json
import logging
import asyncio
from typing import Dict, Any, Optional, Set, Callable
from datetime import datetime

from app.core.json_helpers import _dumps
import pytz

try:
    import websockets
    from websockets.client import WebSocketClientProtocol
except ImportError:
    websockets = None
    WebSocketClientProtocol = None

from app.core.config import settings
from app.realtime.event_handler import event_handler, EventType

logger = logging.getLogger(__name__)


class KISWebSocketClient:
    """
    KIS API WebSocket 클라이언트
    
    실시간 호가/체결 데이터 수신
    """
    
    def __init__(self, app_key: str, app_secret: str):
        if websockets is None:
            raise ImportError("websockets 패키지가 설치되지 않았습니다: pip install websockets")
        
        self.app_key = app_key
        self.app_secret = app_secret
        self.ws: Optional[WebSocketClientProtocol] = None
        self.subscriptions: Set[str] = set()  # 구독 중인 종목코드
        self.is_connected = False
        self.kst = pytz.timezone("Asia/Seoul")
        
        # WebSocket URL (모의투자 vs 실계좌)
        if "vts" in settings.KOREA_INVESTMENT_API_BASE_URL.lower():
            self.ws_url = "ws://ops.koreainvestment.com:31000"  # 모의투자
        else:
            self.ws_url = "ws://ops.koreainvestment.com:21000"  # 실계좌
        
        logger.info(f"🔌 WebSocket 클라이언트 초기화: {self.ws_url}")
    
    async def connect(self):
        """
WebSocket 연결"""
        try:
            logger.info(f"🔌 WebSocket 연결 시도: {self.ws_url}")
            self.ws = await websockets.connect(self.ws_url)
            self.is_connected = True
            logger.info("✅ WebSocket 연결 성공")
            
            await event_handler.emit(EventType.CONNECTED, {"timestamp": datetime.now(self.kst)})
        
        except Exception as e:
            logger.error(f"❌ WebSocket 연결 실패: {e}", exc_info=True)
            self.is_connected = False
            raise
    
    async def disconnect(self):
        """
WebSocket 연결 해제"""
        if self.ws:
            await self.ws.close()
            self.is_connected = False
            logger.info("🔌 WebSocket 연결 해제")
            
            await event_handler.emit(EventType.DISCONNECTED, {"timestamp": datetime.now(self.kst)})
    
    async def subscribe_price(self, stock_code: str):
        """
        주식 현재가 구독
        
        Args:
            stock_code: 종목코드 (6자리)
        """
        if not self.is_connected:
            raise ConnectionError("WebSocket이 연결되지 않았습니다")
        
        # KIS WebSocket 프로토콜
        subscribe_msg = {
            "header": {
                "approval_key": self.app_key,
                "custtype": "P",
                "tr_type": "1",  # 등록
                "content-type": "utf-8"
            },
            "body": {
                "input": {
                    "tr_id": "H0STCNT0",  # 현재가 호가
                    "tr_key": stock_code
                }
            }
        }
        
        await self.ws.send(_dumps(subscribe_msg))
        self.subscriptions.add(stock_code)
        logger.info(f"📡 구독 등록: {stock_code}")
    
    async def unsubscribe_price(self, stock_code: str):
        """주식 현재가 구독 해제"""
        if not self.is_connected:
            return
        
        unsubscribe_msg = {
            "header": {
                "approval_key": self.app_key,
                "custtype": "P",
                "tr_type": "2",  # 해제
                "content-type": "utf-8"
            },
            "body": {
                "input": {
                    "tr_id": "H0STCNT0",
                    "tr_key": stock_code
                }
            }
        }
        
        await self.ws.send(_dumps(unsubscribe_msg))
        self.subscriptions.discard(stock_code)
        logger.info(f"❌ 구독 해제: {stock_code}")
    
    async def listen(self):
        """
        메시지 수신 루프 (무한 대기)
        
        수신한 데이터를 파싱하여 이벤트 발생
        """
        if not self.is_connected:
            raise ConnectionError("WebSocket이 연결되지 않았습니다")
        
        logger.info("👂 WebSocket 메시지 수신 시작")
        
        try:
            async for message in self.ws:
                try:
                    await self._handle_message(message)
                except Exception as e:
                    logger.error(f"❌ 메시지 처리 오류: {e}", exc_info=True)
                    await event_handler.emit(EventType.ERROR, {"error": str(e)})
        
        except websockets.exceptions.ConnectionClosed:
            logger.warning("⚠️ WebSocket 연결 끔김")
            self.is_connected = False
            await event_handler.emit(EventType.DISCONNECTED, {"reason": "connection_closed"})
    
    async def _handle_message(self, message: str):
        """
        수신한 메시지 파싱 및 이벤트 발생
        
        Args:
            message: WebSocket 메시지 (JSON 문자열)
        """
        try:
            data = json.loads(message)
        except json.JSONDecodeError:
            # 텍스트 형식일 수 있음 (KIS API 특성)
            logger.debug(f"텍스트 메시지 수신: {message[:100]}...")
            return
        
        # TR_ID 확인
        tr_id = data.get("header", {}).get("tr_id")
        
        if tr_id == "H0STCNT0":  # 현재가 호가
            await self._handle_price_update(data)
        
        elif tr_id == "H0STCNI0":  # 체결
            await self._handle_execution(data)
        
        else:
            logger.debug(f"알 수 없는 TR_ID: {tr_id}")
    
    async def _handle_price_update(self, data: Dict[str, Any]):
        """현재가 업데이트 처리"""
        body = data.get("body", {})
        output = body.get("output", {})
        
        event_data = {
            "stock_code": output.get("MKSC_SHRN_ISCD", ""),
            "stock_name": output.get("HTS_KOR_ISNM", ""),
            "current_price": int(output.get("STCK_PRPR", 0)),
            "change_rate": float(output.get("PRDY_CTRT", 0)),
            "volume": int(output.get("ACML_VOL", 0)),
            "timestamp": datetime.now(self.kst),
        }
        
        await event_handler.emit(EventType.PRICE_UPDATE, event_data)
    
    async def _handle_execution(self, data: Dict[str, Any]):
        """체결 데이터 처리"""
        body = data.get("body", {})
        output = body.get("output", {})
        
        event_data = {
            "order_id": output.get("ODNO", ""),
            "stock_code": output.get("PDNO", ""),
            "executed_price": int(output.get("CNTR_UNPR", 0)),
            "executed_quantity": int(output.get("CNTR_QTY", 0)),
            "timestamp": datetime.now(self.kst),
        }
        
        await event_handler.emit(EventType.EXECUTION, event_data)
    
    def get_subscriptions(self) -> Set[str]:
        """현재 구독 중인 종목 목록"""
        return self.subscriptions.copy()
    
    async def run(self):
        """
        WebSocket 클라이언트 실행 (연결 + 수신)
        
        자동 재연결 지원
        """
        while True:
            try:
                await self.connect()
                await self.listen()
            
            except Exception as e:
                logger.error(f"❌ WebSocket 오류: {e}")
                await event_handler.emit(EventType.ERROR, {"error": str(e)})
                
                # 5초 후 재연결 시도
                logger.info("🔄 5초 후 재연결 시도...")
                await asyncio.sleep(5)


# 클라이언트 인스턴스 생성 헬퍼
def create_websocket_client() -> KISWebSocketClient:
    """
KIS WebSocket 클라이언트 인스턴스 생성"""
    return KISWebSocketClient(
        app_key=settings.KOREA_INVESTMENT_APP_KEY,
        app_secret=settings.KOREA_INVESTMENT_APP_SECRET
    )