#!/usr/bin/env python3
"""키움 실시간 시세 허브 (WebSocket) — A단계 단독 모듈.

조회 전용 REST(kiwoom_client)와 별개로, 키움 실시간 시세 WebSocket을 상시 청취해
최신 현재가(0B 주식체결)·VI 발동/해제(1h)를 메모리에 들고 있다가 읽어주는 컴포넌트.

behive_web 프로세스 안의 데몬 스레드로 띄울 예정이지만(B단계), 이 파일 자체는
behive_web을 전혀 import/수정하지 않는다. 단독 실행으로 실접속·파싱을 먼저 검증한다.

  python3 realtime_hub.py test 005930 000660   # 지정 종목 실시간 수신을 콘솔에 출력

설계:
  - 연결 1개(시세는 계좌 무관 → 토큰 아무 계좌나 재사용)
  - 흐름: connect → LOGIN → REG(구독) → 수신 루프(PING 에코·REAL 파싱) → 끊기면 백오프 재연결+재구독
  - 메모리 상태 저장소(lock): quotes{code:{...}}, vi{code:{...}}
  - hub가 콜드/끊김이어도 호출측은 REST로 폴백 가능하도록, 읽기 메서드는 없으면 None 반환

키움 WS 프로토콜:
  - 엔드포인트: wss://api.kiwoom.com:10000/api/dostk/websocket (운영)
  - LOGIN/PING 핸드셰이크는 공식 PDF에 표가 없어 키움 공식 샘플 규약을 따름:
      LOGIN 송신: {"trnm":"LOGIN","token":"<access_token>"}  (Bearer 접두어 없이 raw 토큰)
      서버가 {"trnm":"PING"} 주기 송신 → 같은 메시지 그대로 에코
  - REG/REAL 스펙·필드코드는 PDF(국내주식>실시간시세) 기준:
      0B 주식체결: 10=현재가 11=전일대비 12=등락율 20=체결시간 16/17/18=시고저 9081=거래소
      1h VI발동/해제: 9068=발동구분 1221=발동가격 1224=해제시각 1225=적용구분 1489=등락율 9069=방향
"""
from __future__ import annotations

import json
import sys
import threading
import time
from datetime import datetime, timezone, timedelta

import websocket  # websocket-client (동기식). 기존 threading 서버와 궁합.

import kiwoom_client as kc

KST = timezone(timedelta(hours=9))

# 재연결 백오프 (초) — 점증, 상한 30s
_BACKOFF_START = 2.0
_BACKOFF_MAX = 30.0
# 토큰 만료 전 선제 재연결 (초). 키움 토큰 ~30분 → 25분마다 새 연결로 갈아탐.
_MAX_CONN_AGE = 25 * 60
# recv 타임아웃 — 이 시간 내 아무것도 안 오면 루프 돌며 stop/age 체크
_RECV_TIMEOUT = 5.0
# VI 배지 자동 만료 (초) — 해제 1h 메시지 미수신 대비. VI 단일가 2분 + 임의연장 ≤30초 + 여유 30초.
_VI_ACTIVE_TTL = 180.0
# 장 마감 시간대 대기 — 개장 임박(평일 07:00~08:00)엔 1분, 그 외엔 다음 영업일 07:00까지 자되 6시간 상한.
_PREOPEN_POLL_SEC = 60.0       # 개장 임박: 1분마다 → 08:00 개장을 1분 안에 포착
_CLOSED_WAIT_CAP = 6 * 3600.0  # 밤·주말·휴장: 다음 개장까지 자되 이 상한(자가교정용 — 날짜 계산 오류 시에도 6h마다 재확인)


def _ws_url() -> str:
    """REST base_url에서 WS 엔드포인트 도출. https→wss, :10000, /api/dostk/websocket."""
    base = kc.base_url()  # 예: https://api.kiwoom.com 또는 https://mockapi.kiwoom.com
    host = base.split('://', 1)[-1].rstrip('/')
    return f'wss://{host}:10000/api/dostk/websocket'


def _abs_int(s) -> int:
    """부호 포함 문자열 → 절대값 int (현재가·호가용)."""
    return abs(kc._to_int(s))


class RealtimeHub:
    """키움 실시간 시세 허브. start()로 데몬 스레드 기동, get_quote/get_vi로 읽기."""

    def __init__(self, account_label: str | None = None):
        self._label = account_label or kc._default_account_label()
        self._lock = threading.Lock()
        self._quotes: dict[str, dict] = {}   # code -> {price, change, pct, open, high, low, ts, exchange}
        self._vi: dict[str, dict] = {}        # code -> {active, kind, trigger_price, ...}
        self._subs: set[str] = set()          # 현재 구독 종목코드(정규화)
        self._ws: websocket.WebSocket | None = None
        self._ws_lock = threading.Lock()      # send 직렬화
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._connected = threading.Event()
        self.last_error: str | None = None
        self.connected_at: float = 0.0

    # ---- 공개 읽기 API (없으면 None → 호출측 REST 폴백) ----
    def get_quote(self, code: str) -> dict | None:
        with self._lock:
            q = self._quotes.get(kc._clean_code(code))
            return dict(q) if q else None

    def get_quotes(self, codes: list[str]) -> dict[str, dict]:
        with self._lock:
            out = {}
            for c in codes:
                cc = kc._clean_code(c)
                if cc in self._quotes:
                    out[cc] = dict(self._quotes[cc])
            return out

    def get_vi(self, code: str) -> dict | None:
        with self._lock:
            v = self._vi.get(kc._clean_code(code))
            return self._vi_expired(v) if v else None

    def active_vis(self) -> dict[str, dict]:
        with self._lock:
            out = {}
            for c, v in self._vi.items():
                ev = self._vi_expired(v)
                if ev.get('active'):
                    out[c] = ev
            return out

    @staticmethod
    def _vi_expired(v: dict) -> dict:
        """VI 자동 만료 — 해제(1h) 메시지가 안 오는 경우 대비.
        VI 단일가는 2분 + 임의연장 ≤30초라, 발동 후 _VI_ACTIVE_TTL 초과 시 읽기 시점에 active=False.
        저장 상태는 건드리지 않는다(실증·디버깅용 원본 유지)."""
        ev = dict(v)
        if ev.get('active') and ev.get('ts'):
            try:
                age = (datetime.now(KST) - datetime.fromisoformat(ev['ts'])).total_seconds()
                if age > _VI_ACTIVE_TTL:
                    ev['active'] = False
            except Exception:
                pass
        return ev

    def is_connected(self) -> bool:
        return self._connected.is_set()

    # ---- 구독 관리 ----
    def set_subscriptions(self, codes: list[str]) -> None:
        """구독 종목 목록을 통째로 교체. 연결돼 있으면 즉시 REG/REMOVE 반영."""
        new = {kc._clean_code(c) for c in codes if c}
        with self._lock:
            add = new - self._subs
            rem = self._subs - new
            self._subs = new
        if self._connected.is_set():
            if add:
                self._send_reg(sorted(add))
            if rem:
                self._send_remove(sorted(rem))

    # ---- 생명주기 ----
    def start(self) -> None:
        if self._thread and self._thread.is_alive():
            return
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, name='realtime-hub', daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        with self._ws_lock:
            if self._ws:
                try:
                    self._ws.close()
                except Exception:
                    pass

    # ---- 내부: 연결 루프 ----
    def _market_open(self) -> bool:
        """장 시간대(정규장·NXT) 여부. behive_web._market_phase_state 재사용(지연 import로 순환 회피).
        판정 실패 시 보수적으로 True(연결 유지) — 게이팅 때문에 시세가 끊기는 일은 없게."""
        try:
            import behive_web
            return bool(behive_web._market_phase_state().get('active'))
        except Exception:
            return True

    @staticmethod
    def _next_preopen_dt(now: datetime, holidays: set) -> datetime:
        """now 이후 가장 가까운 영업일(평일·비휴장) 07:00 datetime. 주말·휴장·연휴 건너뜀."""
        for i in range(0, 15):  # 최대 2주(연휴) 탐색
            d = now + timedelta(days=i)
            if d.weekday() < 5 and d.strftime('%Y-%m-%d') not in holidays:
                cand = d.replace(hour=7, minute=0, second=0, microsecond=0)
                if cand > now:
                    return cand
        return now + timedelta(hours=6)  # 안전 fallback

    def _closed_wait_seconds(self) -> float:
        """장 외 대기 시간.
        - 평일 장외(closed)이고 NXT 프리마켓 임박(07:00~08:00) → 60s (촘촘히)
        - 그 외(밤·주말·휴장) → 다음 영업일 07:00까지 자되 6시간 상한 (헛바퀴 제거 + 자가교정).
        판정 실패 시 보수적으로 60s."""
        try:
            import behive_web
            now = datetime.now(KST)
            hm = now.hour * 60 + now.minute
            if behive_web._market_phase_state().get('phase') == 'closed' and 7 * 60 <= hm < 8 * 60:
                return _PREOPEN_POLL_SEC
            target = self._next_preopen_dt(now, behive_web._load_holidays())
            return max(60.0, min((target - now).total_seconds(), _CLOSED_WAIT_CAP))
        except Exception:
            return _PREOPEN_POLL_SEC

    def _run(self) -> None:
        backoff = _BACKOFF_START
        while not self._stop.is_set():
            if not self._market_open():
                # 장 외 → 연결 안 하고 대기. 개장 임박엔 촘촘히, 밤·주말·휴장엔 드물게(헛바퀴 최소화).
                self._connected.clear()
                self._stop.wait(self._closed_wait_seconds())
                continue
            try:
                self._connect_once()
                backoff = _BACKOFF_START  # 정상 종료(선제 재연결·장 마감 등) → 백오프 리셋
            except Exception as e:
                self.last_error = f'{type(e).__name__}: {e}'
                self._connected.clear()
                sys.stderr.write(f'[realtime-hub] 연결 끊김/실패: {self.last_error}\n')
                if self._stop.is_set():
                    break
                self._stop.wait(backoff)
                backoff = min(backoff * 2, _BACKOFF_MAX)

    def _connect_once(self) -> None:
        url = _ws_url()
        token = kc.issue_token(self._label)  # 매 연결마다 fresh (캐시·만료는 issue_token이 관리)
        ws = websocket.create_connection(url, timeout=10)
        ws.settimeout(_RECV_TIMEOUT)
        with self._ws_lock:
            self._ws = ws
        try:
            # 1) LOGIN
            ws.send(json.dumps({'trnm': 'LOGIN', 'token': token}))
            login = json.loads(ws.recv())
            if login.get('trnm') == 'LOGIN' and login.get('return_code') not in (0, '0', None):
                raise RuntimeError(f'LOGIN 실패: {login.get("return_msg")} ({login})')
            self._connected.set()
            self.connected_at = time.time()
            self.last_error = None
            sys.stderr.write(f'[realtime-hub] LOGIN 성공 @ {url}\n')

            # 2) 현재 구독 종목 REG
            with self._lock:
                codes = sorted(self._subs)
            if codes:
                self._send_reg(codes)

            # 3) 수신 루프
            _last_mkt_check = time.time()
            empty_streak = 0  # 연속 빈응답 카운터 — tight busy-loop 방지
            while not self._stop.is_set():
                if time.time() - self.connected_at > _MAX_CONN_AGE:
                    sys.stderr.write('[realtime-hub] 토큰 만료 전 선제 재연결\n')
                    return  # 정상 종료 → _run이 즉시 재연결
                # 장 마감되면 연결 종료 → _run이 대기 모드로 전환. 30초마다만 확인(매 recv마다 X).
                if time.time() - _last_mkt_check > 30:
                    _last_mkt_check = time.time()
                    if not self._market_open():
                        sys.stderr.write('[realtime-hub] 장 마감 → 연결 종료, 대기 모드\n')
                        return
                try:
                    raw = ws.recv()
                except websocket.WebSocketTimeoutException:
                    empty_streak = 0  # 타임아웃은 정상 idle
                    continue
                if not raw:
                    # 빈응답이 예외 없이 연속되면(드묾) tight loop → 5회 누적 시 재연결로 탈출
                    empty_streak += 1
                    if empty_streak >= 5:
                        sys.stderr.write('[realtime-hub] 빈 응답 연속 → 연결 재설정\n')
                        return
                    continue
                empty_streak = 0
                self._handle(raw)
        finally:
            self._connected.clear()
            with self._ws_lock:
                self._ws = None
            try:
                ws.close()
            except Exception:
                pass

    # ---- 송신 ----
    def _send(self, payload: dict) -> None:
        with self._ws_lock:
            if self._ws is None:
                return
            self._ws.send(json.dumps(payload))

    @staticmethod
    def _al(codes: list[str]) -> list[str]:
        # 통합(SOR) 거래소코드 `_AL` 접미사 — 정규장 KRX 체결 + NXT 시간대 NXT 체결을
        # 한 구독으로 수신. bare 코드는 KRX 단독이라 NXT(15:30~20:00)에 시세가 끊긴다. (PDF 0B/1h 스펙)
        return [f'{c}_AL' for c in codes]

    def _send_reg(self, codes: list[str]) -> None:
        # 0B(체결) + 1h(VI) 동시 등록. refresh=1 → 기존 구독 유지하며 추가.
        self._send({
            'trnm': 'REG', 'grp_no': '1', 'refresh': '1',
            'data': [{'item': self._al(codes), 'type': ['0B', '1h']}],
        })

    def _send_remove(self, codes: list[str]) -> None:
        self._send({
            'trnm': 'REMOVE', 'grp_no': '1',
            'data': [{'item': self._al(codes), 'type': ['0B', '1h']}],
        })

    # ---- 수신 처리 ----
    def _handle(self, raw: str) -> None:
        try:
            msg = json.loads(raw)
        except (ValueError, TypeError):
            return
        trnm = msg.get('trnm')
        if trnm == 'PING':
            # PING은 받은 그대로 에코 (키움 keepalive 규약)
            with self._ws_lock:
                if self._ws is not None:
                    self._ws.send(raw)
            return
        if trnm == 'REAL':
            for blk in msg.get('data') or []:
                self._handle_real(blk)
            return
        if trnm in ('REG', 'REMOVE'):
            rc = msg.get('return_code')
            if rc not in (0, '0', None):
                sys.stderr.write(f'[realtime-hub] {trnm} 거부: rc={rc} msg={msg.get("return_msg")}\n')
            else:
                sys.stderr.write(f'[realtime-hub] {trnm} ok (rc={rc})\n')
            return
        # LOGIN 응답 등 기타는 무시

    def _handle_real(self, blk: dict) -> None:
        typ = blk.get('type')
        # 통합(_AL)·NXT(_NX) 구독은 item이 `039490_AL` 형태로 echo됨 → 거래소 접미사를 떼고
        # bare 코드로 저장해야 get_quotes(bare) 매칭이 유지된다. 키움 종목코드엔 '_'가 없다.
        code = kc._clean_code((blk.get('item') or '').split('_')[0])
        vals = blk.get('values') or {}
        if not code:
            return
        if typ == '0B':  # 주식체결 → 현재가
            with self._lock:
                self._quotes[code] = {
                    'price': _abs_int(vals.get('10')),
                    'change': kc._to_int(vals.get('11')),
                    'pct': kc._to_float(vals.get('12')),
                    'open': _abs_int(vals.get('16')),   # 시가
                    'high': _abs_int(vals.get('17')),   # 고가
                    'low': _abs_int(vals.get('18')),    # 저가
                    'exchange': (vals.get('9081') or '').strip(),
                    'ts': datetime.now(KST).isoformat(),
                }
        elif typ == '1h':  # VI 발동/해제
            # active 판정: 9068 VI발동구분. 해제시각(1224) 유무와 함께 실증 확인 예정(A단계).
            # 실증 로그 — VI는 드물어 부담 없음. 2026-06-11 관찰: 발동 후 해제 1h가 안 오고
            # 1224(해제시각)가 발동시각과 동일 → 읽기 시점 TTL 만료(_vi_expired)로 방어 중.
            print(f'[1h raw] {code} {vals}', file=sys.stderr, flush=True)
            kind_raw = (vals.get('9068') or '').strip()
            with self._lock:
                self._vi[code] = {
                    'active': kind_raw not in ('', '0', '2'),  # ⚠️ 매핑 실증 확인 필요
                    'vi_gubun': kind_raw,
                    'apply_kind': (vals.get('1225') or '').strip(),  # 정적/동적/동적+정적
                    'trigger_price': _abs_int(vals.get('1221')),
                    'trigger_pct': kc._to_float(vals.get('1489')),
                    'release_time': (vals.get('1224') or '').strip(),
                    'direction': (vals.get('9069') or '').strip(),
                    'count': kc._to_int(vals.get('1490')),  # 당일 VI 발동 횟수
                    'ts': datetime.now(KST).isoformat(),
                }


# ---------------- CLI (A단계 단독 검증) ----------------
def _cmd_test(codes: list[str]) -> int:
    if not codes:
        codes = ['005930', '000660']  # 삼성전자, SK하이닉스
    print(f'[test] WS URL = {_ws_url()}')
    print(f'[test] 구독 종목 = {codes}  (Ctrl-C 종료)')
    hub = RealtimeHub()
    hub.set_subscriptions(codes)
    hub.start()
    last_print: dict[str, str] = {}
    try:
        while True:
            time.sleep(1)
            if not hub.is_connected():
                continue
            for c in codes:
                cc = kc._clean_code(c)
                q = hub.get_quote(cc)
                if q:
                    line = f"{q['price']:,} ({q['pct']:+.2f}%) {q['exchange']}"
                    if last_print.get(cc) != line:
                        print(f'  [{cc}] {line}  @ {q["ts"][11:19]}')
                        last_print[cc] = line
                v = hub.get_vi(cc)
                if v and v.get('active'):
                    vk = f"VI:{cc} {v['apply_kind']} 발동가 {v['trigger_price']:,} ({v['trigger_pct']:+.2f}%)"
                    if last_print.get(f'vi_{cc}') != vk:
                        print(f'  ⚡ {vk}')
                        last_print[f'vi_{cc}'] = vk
    except KeyboardInterrupt:
        print('\n[test] 종료')
        hub.stop()
    return 0


def main(argv: list[str]) -> int:
    if not argv or argv[0] != 'test':
        print('usage: python3 realtime_hub.py test [code ...]', file=sys.stderr)
        return 2
    return _cmd_test(argv[1:])


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))