import asyncio
import json
import websockets
import logging
from typing import Dict, Set, Callable, Optional

from app.services.kis_auth import kis_auth
from app.core.crypto import aes_cbc_base64_dec
from app.db.database import SessionLocal
# from app.db.crud import update_stock_price # TODO: Implement CRUD

logger = logging.getLogger(__name__)

class RealtimeManager:
    """
    Manages KIS WebSocket Connection.
    Handles: Connection, Subscription, Decryption, PINGPONG.
    """
    WS_URL_REAL = "ws://ops.koreainvestment.com:21000"
    
    def __init__(self):
        self.ws: Optional[websockets.WebSocketClientProtocol] = None
        self.approval_key: Optional[str] = None
        self.subscribed_codes: Set[str] = set()
        self.running = False
        self.data_map: Dict[str, Dict] = {} # Store IV/Key for encrypted TRs

    async def start(self):
        """
        Main loop: Connect -> Authenticate -> Listen
        """
        self.running = True
        while self.running:
            try:
                # 1. Get Approval Key
                self.approval_key = await kis_auth.get_approval_key()
                
                logger.info(f"Connecting to KIS WS: {self.WS_URL_REAL}")
                async with websockets.connect(self.WS_URL_REAL, ping_interval=None) as websocket:
                    self.ws = websocket
                    logger.info("Connected.")
                    
                    # 2. Resubscribe if recovering connection
                    if self.subscribed_codes:
                        await self._resubscribe_all()

                    # 3. Listen Loop
                    await self._listen()
                    
            except Exception as e:
                logger.error(f"WS Connection Error: {e}. Retrying in 5s...")
                await asyncio.sleep(5)

    async def stop(self):
        self.running = False
        if self.ws:
            await self.ws.close()

    async def subscribe(self, stock_code: str, type="price"):
        """
        Subscribe to a stock.
        type: 'price' (H0STCNT0 - 체결가)
        """
        if not self.ws or not self.approval_key:
            logger.warning("WS not ready. Adding to pending list.")
            self.subscribed_codes.add(stock_code)
            return

        # Domestic Realtime Price TR ID: H0STCNT0
        tr_id = "H0STCNT0" 
        tr_key = stock_code
        
        payload = {
            "header": {
                "approval_key": self.approval_key,
                "custtype": "P",
                "tr_type": "1", # 1=Register, 2=Unregister
                "content-type": "utf-8"
            },
            "body": {
                "input": {
                    "tr_id": tr_id,
                    "tr_key": tr_key
                }
            }
        }
        
        await self.ws.send(json.dumps(payload))
        self.subscribed_codes.add(stock_code)
        logger.info(f"Subscribed to {stock_code}")

    async def _resubscribe_all(self):
        for code in self.subscribed_codes:
            await self.subscribe(code)

    async def _listen(self):
        async for message in self.ws:
            try:
                # Message can be plain text provided by library, or bytes
                if isinstance(message, bytes):
                    message = message.decode('utf-8')

                # KIS sends data in specific formats.
                # 1. JSON (Control Messages, PINGPONG, Subscription Ack)
                # 2. Text/Pipe separated (Real Data) - Usually starts with 0 or 1
                
                first_char = message[0]
                
                if first_char in ['{', '[']: 
                    # JSON Message
                    data = json.loads(message)
                    header = data.get('header', {})
                    tr_id = header.get('tr_id')

                    if tr_id == "PINGPONG":
                        await self.ws.send(message) # Echo back
                        logger.debug("PINGPONG handled")
                        
                    elif 'body' in data:
                        # Subscription Ack
                        # Store IV/Key if encryption is enabled (msg1 often contains 'ENCRYPT')
                        # But for Brokerage API, H0STCNT0 is usually plaintext unless configured otherwise.
                        # If encrypted, 'iv' and 'key' are in body['output']
                        pass
                        
                elif first_char in ['0', '1']:
                    # Real Data: 0|TR_ID|DATA_CNT|DATA...
                    parts = message.split('|')
                    if len(parts) < 4:
                        continue
                        
                    tr_id = parts[1]
                    raw_data = parts[3]
                    
                    # Decryption Check
                    # If this tr_id was registered as encrypted, decrypt it.
                    # For now assuming Plaintext for H0STCNT0 as per standard Personal API.
                    
                    # Parse Data
                    if tr_id == "H0STCNT0": # Domestic Price
                        # Data format: TIME^PRICE^...
                        # We need to look up format spec.
                        # Simple implementation: just log or split
                        fields = raw_data.split('^')
                        if len(fields) > 2:
                            current_price = fields[2] # Example index
                            # TODO: Update DB
                            # print(f"Price Update: {current_price}")
                            pass
                            
            except Exception as e:
                logger.error(f"Error processing WS message: {e}")

realtime_manager = RealtimeManager()