KisStock0/backend/websocket_manager.py

import asyncio
import websockets
import json
import logging
import datetime
from typing import List, Set
from database import SessionLocal, Stock, StockPrice
from kis_api import kis

logger = logging.getLogger("WEBSOCKET")

class KisWebSocketManager:
    def __init__(self):
        self.active_frontend_connections: List[any] = []
        self.subscribed_codes: Set[str] = set()
        self.running = False
        self.approval_key = None
        self.msg_queue = asyncio.Queue() # For outgoing subscription requests
        
        # KIS Environment
        if kis.is_paper:
            self.url = "ws://ops.koreainvestment.com:31000"
        else:
            self.url = "ws://ops.koreainvestment.com:21000"

    # ... (connect/disconnect/broadcast remains same)

    # ... (connect/disconnect/broadcast remains same)

    # ... (_handle_realtime_data remains same)

    async def subscribe_stock(self, code):
        if code in self.subscribed_codes:
            return
        self.subscribed_codes.add(code)
        await self.msg_queue.put(code)
        logger.info(f"Queued Subscription for {code}")
            
    async def connect_frontend(self, websocket):
        await websocket.accept()
        self.active_frontend_connections.append(websocket)
        logger.info(f"Frontend Client Connected. Total: {len(self.active_frontend_connections)}")

    def disconnect_frontend(self, websocket):
        if websocket in self.active_frontend_connections:
            self.active_frontend_connections.remove(websocket)
            logger.info("Frontend Client Disconnected")

    async def broadcast_to_frontend(self, message: dict):
        # Broadcast to all connected frontend clients
        for connection in self.active_frontend_connections:
            try:
                await connection.send_json(message)
            except Exception as e:
                logger.error(f"Broadcast error: {e}")
                self.disconnect_frontend(connection)

    async def start_kis_socket(self):
        self.running = True
        logger.info(f"Starting KIS WebSocket Service... Target: {self.url}")
        
        while self.running:
            # 1. Ensure Approval Key
            if not self.approval_key:
                self.approval_key = kis.get_websocket_key()
                if not self.approval_key:
                    logger.error("Failed to get WebSocket Approval Key. Retrying in 10s...")
                    await asyncio.sleep(10)
                    continue
                logger.info(f"Got WS Key: {self.approval_key[:10]}...")

            # 2. Connect
            try:
                # KIS doesn't use standard ping frames often, handle manually or disable auto-ping
                async with websockets.connect(self.url, ping_interval=None, open_timeout=20) as ws:
                    logger.info("Connected to KIS WebSocket Server")
                    
                    # Process initial subscriptions
                    for code in self.subscribed_codes:
                        await self._send_subscription(ws, code)
                        
                    while self.running:
                        try:
                            # 1. Check for incoming data
                            msg = await asyncio.wait_for(ws.recv(), timeout=0.1)
                            
                            # PING/PONG (String starting with 0 or 1 usually means data)
                            if msg[0] in ['0', '1']: 
                                await self._handle_realtime_data(msg)
                            else:
                                # JSON Message (System, PINGPONG)
                                try:
                                    data = json.loads(msg)
                                    if data.get('header', {}).get('tr_id') == 'PINGPONG':
                                        await ws.send(msg) # Echo back
                                        continue
                                except:
                                    pass 
                                    
                        except asyncio.TimeoutError:
                            pass
                            
                        except websockets.ConnectionClosed:
                            logger.warning("KIS WS Closed. Reconnecting...")
                            break
                            
            except Exception as e:
                logger.error(f"WS Connection Error: {e}")
                # If auth failed (maybe expired key?), clear key to force refresh
                # simplified check: if "Approval key" error in exception message?
                # For now just retry.
                await asyncio.sleep(5)

    async def _handle_realtime_data(self, msg: str):
        # Format: 0|TR_ID|DATA_CNT|Code^Time^Price...
        try:
            parts = msg.split('|')
            if len(parts) < 4: return
            
            tr_id = parts[1]
            data_part = parts[3]
            
            if tr_id == "H0STCNT0": # Domestic Stock Price
                # Data format: Code^Time^CurrentPrice^Sign^Change...
                # Actually, data_part is delimiter separated.
                values = data_part.split('^')
                code = values[0]
                price = values[2]
                change = values[4]
                rate = values[5]
                
                # Broadcast
                payload = {
                    "type": "PRICE",
                    "code": code,
                    "price": price,
                    "change": change,
                    "rate": rate,
                    "timestamp": datetime.datetime.now().isoformat()
                }
                await self.broadcast_to_frontend(payload)
                
                # Update DB (Optional? Too frequent writes maybe bad)
                # Let's save only significant updates or throttle?
                # For now just log/broadcast.
                
        except Exception as e:
            logger.error(f"Data Parse Error: {e} | Msg: {msg[:50]}")

    async def subscribe_stock(self, code):
        if code in self.subscribed_codes:
            return
        self.subscribed_codes.add(code)
        # If socket is active, send subscription (Implementation complexity: need access to active 'ws' object)
        # Will handle by restarting connection or using a queue? 
        # Better: just set it in set, and the main loop will pick it up on reconnect, 
        # BUT for immediate sub, we need a way to signal the running loop. 
        # For MVP, let's assume we subscribe on startup or bulk. 
        # Real-time dynamic sub needs a queue.
        logger.info(f"Subscribed to {code} (Pending next reconnect/sweep)")
        
    async def _send_subscription(self, ws, code):
        # Domestic Stock Realtime Price: H0STCNT0
        body = {
            "header": {
                "approval_key": self.approval_key,
                "custtype": "P",
                "tr_type": "1", # 1: Register, 2: Unregister
                "content-type": "utf-8"
            },
            "body": {
                "input": {
                    "tr_id": "H0STCNT0", 
                    "tr_key": code 
                }
            }
        }
        await ws.send(json.dumps(body))
        logger.info(f"Sent Subscription Request for {code}")

ws_manager = KisWebSocketManager()