"""
kis_trader/network/ws_manager.py — 단일 WebSocket 허브 (Event Bus)
====================================================================
설계 목적:
  * 두 전략(스캘핑/꼬리잡기)이 각자 WS 연결을 띄우면 같은 종목 2번 구독 → 토큰/approval_key
    경합 + 계정 차단 위험 → 프로세스 전체에서 **단 하나의 WS 세션**만 띄운다.
  * 기존 ``kis_ws.KISWebSocketPriceCache`` + ``CandleAggregator`` 를 그대로 재사용.
  * 전략별 "구독 관심 종목" 을 **레퍼런스 카운팅**으로 관리. 한 전략이 구독 해제해도
    다른 전략이 구독 중이면 WS 에서 해제되지 않는다.

공용 API:
  - start() / stop()
  - subscribe(code, owner) / unsubscribe(code, owner)
  - sync_targets(owner, codes) : 한 전략의 관심 종목을 한 번에 동기화
  - get_price(code) / get_candles(code, tf, n)
  - fill_gap(codes=None) : 갭 보정 (WS 연결 직후 자동 + 수동 호출)
"""
from __future__ import annotations

import queue
import random
import threading
import time
from collections import defaultdict
from typing import Dict, Iterable, List, Optional, Set

from ..utils.env import get_env_bool, get_env_from_db, get_env_int
from ..utils.logger import get_logger

logger = get_logger("kis_trader.ws")

# 기존 kis_ws 모듈 재사용 (검증된 로직 보존 원칙)
# 역할 분리 정책 (kis_scalping_ver2 / kis_short_ver3 와 동일):
#   - 실시간 시세(WS)         : KIS 실전키     (is_mock=False 고정)
#   - 매수/매도 주문·계좌      : KIS (mock 여부는 KIS_MOCK)
#   - 유니버스(거래량 순위)    : KIS REST volume-rank
#   - 과거 봉 웜업·갭보정      : 키움 ka10080 (1/3/15/60분 native 지원)
#                                → 키움 키 없으면 KIS 1분봉 fallback
try:
    from kis_ws import (
        CandleAggregator,
        KISWebSocketPriceCache,
        _get_kiwoom_creds,
        get_kiwoom_candles_df,
    )
    _KIS_WS_AVAILABLE = True
except ImportError as _e:
    _KIS_WS_AVAILABLE = False
    # 서브심볼 import 실패 대응 (패키지만 있고 키움 함수 없는 구버전)
    try:
        from kis_ws import CandleAggregator, KISWebSocketPriceCache  # type: ignore
        _KIS_WS_AVAILABLE = True
        _get_kiwoom_creds = None  # type: ignore[assignment]
        get_kiwoom_candles_df = None  # type: ignore[assignment]
        logger.warning(
            "kis_ws 에 키움 함수 없음 → 갭보정 키움 fallback 비활성 "
            "(KIS 1분봉 전용): %s", _e,
        )
    except ImportError as _e2:
        logger.warning("kis_ws 모듈 import 실패 → WS 기능 비활성: %s", _e2)


class WSManager:
    """
    단일 WS 허브.  전략은 이 매니저를 공유하고, subscribe/unsubscribe 시 owner 를 전달한다.

    reference counting 예시:
        subscribe("005930", owner="SCALP")   # refs[005930]={SCALP}       → WS subscribe
        subscribe("005930", owner="SHORT")   # refs[005930]={SCALP,SHORT} → (already subscribed)
        unsubscribe("005930", owner="SCALP") # refs[005930]={SHORT}       → keep
        unsubscribe("005930", owner="SHORT") # refs[005930]=set()         → WS unsubscribe
    """

    def __init__(self, *, db, kis_client):
        self.db = db
        self.kis_client = kis_client

        self.ws_cache: Optional["KISWebSocketPriceCache"] = None
        self.candle_agg: Optional["CandleAggregator"] = None

        # owner(전략ID) → 관심 코드 집합
        self._owner_codes: Dict[str, Set[str]] = defaultdict(set)
        # code → 보유 중인 owner 집합 (ref counting)
        self._code_refs: Dict[str, Set[str]] = defaultdict(set)
        # 영구 구독(시장방향 ETF 등)
        self._permanent_codes: Set[str] = set()
        self._lock = threading.Lock()

        # ── 갭보정 비동기 파이프라인 ─────────────────────────────
        # (전략 쓰레드에서 subscribe() 시 동기 REST 호출하면 매수 체크가
        #  수 분간 블로킹됨 → 백그라운드 워커 큐로 이관)
        self._gap_q: "queue.Queue[str]" = queue.Queue(maxsize=1024)
        self._gap_filled: Set[str] = set()          # 이미 갭보정 완료한 코드
        self._gap_inflight: Set[str] = set()        # 큐에 등록/처리 중인 코드
        self._gap_lock = threading.Lock()
        self._gap_worker_thread: Optional[threading.Thread] = None
        # 전체 재갭보정(재접속 시) 중복 트리거 방지
        self._bulk_refill_running = False

    # ------------------------------------------------------------------
    # 시작/종료
    # ------------------------------------------------------------------
    def start(self) -> bool:
        """WS 세션 시작. 실패 시 False (봇은 REST 폴백으로 동작)."""
        if not _KIS_WS_AVAILABLE:
            logger.warning("kis_ws 미설치 → WS 허브 비활성 (REST 폴백만 동작)")
            return False

        # ── [중요] WS 는 데이터 수신용이므로 무조건 실전 서버로 접속 ──
        # kis_scalping_ver2 와 동일 정책 (모의 계좌라도 시세는 실전 필요)
        ws_app_key = get_env_from_db("KIS_APP_KEY_REAL", "") or self.kis_client.app_key
        ws_app_secret = get_env_from_db("KIS_APP_SECRET_REAL", "") or self.kis_client.app_secret
        if not ws_app_key or not ws_app_secret:
            logger.warning("KIS 실전 키 없음 → WS 허브 비활성")
            return False

        try:
            self.ws_cache = KISWebSocketPriceCache(
                app_key=ws_app_key,
                app_secret=ws_app_secret,
                is_mock=False,  # 시세는 실전 서버 고정
            )
            # 봉 타임프레임: 스캘핑(1분) + 꼬리잡기(3분) + 추세 필터(15/60분)
            # SCALP/SHORT 양쪽 전략이 쓰는 모든 TF 포함
            tfs = self._resolve_timeframes()
            self.candle_agg = CandleAggregator(db=self.db, timeframes=tfs)
            self.ws_cache.attach_candle_aggregator(self.candle_agg)

            ok = self.ws_cache.start()
            if not ok:
                logger.warning("WS start() 실패 → 비활성")
                self.ws_cache = None
                self.candle_agg = None
                return False

            # 갭보정 백그라운드 워커 가동 (subscribe 논블로킹 보장)
            self._start_gap_worker()

            # 영구 구독 (KOSPI/KOSDAQ ETF 등)
            self._load_permanent_codes()
            for code in sorted(self._permanent_codes):
                self.ws_cache.subscribe(code)
                self._enqueue_gap_fill(code)
                logger.info("📡 [영구구독] %s", code)

            # 연결 성공 후 자동 갭 보정 등록 (WS 재접속 시 전체 재갭보정)
            self.ws_cache.set_on_connected_callback(self._trigger_bulk_refill_async)

            logger.info(
                "✅ WSManager 활성 (tfs=%s, permanent=%d, gap_worker=ON)",
                tfs, len(self._permanent_codes),
            )
            return True
        except Exception as e:
            logger.error("WS 초기화 예외: %s", e)
            self.ws_cache = None
            self.candle_agg = None
            return False

    def stop(self) -> None:
        try:
            if self.ws_cache:
                self.ws_cache.stop(clear_subscriptions=True)
        except Exception as e:
            logger.debug("WS stop 실패: %s", e)

    @property
    def is_active(self) -> bool:
        return bool(self.ws_cache and self.ws_cache.is_active)

    # ------------------------------------------------------------------
    # 구독 관리 (Reference Counting)
    # ------------------------------------------------------------------
    def subscribe(self, code: str, owner: str) -> None:
        """
        한 전략(owner)이 해당 종목에 관심 등록.
        ※ 갭보정은 **백그라운드 워커 큐**로 위임하여 전략 쓰레드를 블록하지 않음.
          (예전: 여기서 REST 4개 TF 동기 호출 → 매수 체크 2~3분 지연)
        """
        if not code or not owner:
            return
        with self._lock:
            self._owner_codes[owner].add(code)
            first_ref = not self._code_refs[code]
            self._code_refs[code].add(owner)

        if first_ref and self.ws_cache:
            self.ws_cache.subscribe(code)
            # 신규 구독 → 워커에게 갭보정 위임 (논블로킹)
            self._enqueue_gap_fill(code)

    def unsubscribe(self, code: str, owner: str) -> None:
        """한 전략(owner)이 관심 해제. 다른 전략이 아직 들고 있으면 WS 는 유지."""
        if not code or not owner:
            return
        with self._lock:
            self._owner_codes[owner].discard(code)
            if owner in self._code_refs.get(code, set()):
                self._code_refs[code].discard(owner)
            still_refs = bool(self._code_refs.get(code))
            is_permanent = code in self._permanent_codes

        if not still_refs and not is_permanent and self.ws_cache:
            self.ws_cache.unsubscribe(code)
            if self.candle_agg:
                self.candle_agg.remove_code(code)

    def sync_targets(self, owner: str, codes: Iterable[str]) -> None:
        """
        한 전략의 관심 종목 목록을 통째로 동기화.
        - 기존 관심 종목 중 없어진 것은 unsubscribe
        - 새로 추가된 것은 subscribe
        """
        new_set = {c for c in codes if c}
        with self._lock:
            cur = set(self._owner_codes.get(owner, set()))

        for code in sorted(cur - new_set):
            self.unsubscribe(code, owner)
        for code in sorted(new_set - cur):
            self.subscribe(code, owner)

    # ------------------------------------------------------------------
    # 조회 헬퍼 (전략이 쓰는 API)
    # ------------------------------------------------------------------
    def get_price(self, code: str, max_age_sec: float = 5.0) -> Optional[dict]:
        if not self.ws_cache:
            return None
        try:
            return self.ws_cache.get_price(code, max_age_sec=max_age_sec)
        except Exception:
            return None

    def get_candles(self, code: str, tf: int, n: int = 50) -> list:
        if self.candle_agg:
            try:
                return self.candle_agg.get_candles(code, tf, n)
            except Exception:
                return []
        # CandleAggregator 없으면 DB 폴백
        try:
            return self.db.get_ws_candles(code, tf, limit=n, confirmed_only=True)
        except Exception:
            return []

    def fill_gap(self, codes: Optional[Iterable[str]] = None) -> None:
        """외부에서 수동으로 갭 보정 트리거 (비동기: 큐 등록 후 즉시 리턴)."""
        if codes is None:
            self._trigger_bulk_refill_async()
        else:
            for c in codes:
                self._enqueue_gap_fill(c)

    # ------------------------------------------------------------------
    # 내부: 갭 보정 — 백그라운드 워커 파이프라인
    # ------------------------------------------------------------------
    def _start_gap_worker(self) -> None:
        """갭보정 전담 데몬 워커 스레드 기동 (단일 워커 → REST 레이트리밋 자연 직렬화)."""
        if self._gap_worker_thread and self._gap_worker_thread.is_alive():
            return
        t = threading.Thread(
            target=self._gap_worker_loop,
            name="WS-GapFillWorker",
            daemon=True,
        )
        t.start()
        self._gap_worker_thread = t
        logger.info("✅ 갭보정 워커 스레드 시작 (queue 기반 비동기 처리)")

    def _enqueue_gap_fill(self, code: str) -> None:
        """구독 직후 호출 — 갭보정 큐에 논블로킹 등록.

        중복 방지:
          - 이미 완료(`_gap_filled`) → 스킵
          - 이미 큐/처리 중(`_gap_inflight`) → 스킵
        """
        if not code:
            return
        with self._gap_lock:
            if code in self._gap_filled or code in self._gap_inflight:
                return
            self._gap_inflight.add(code)
        try:
            self._gap_q.put_nowait(code)
        except queue.Full:
            # 큐가 가득 차면 inflight 해제 후 포기 (WS 틱으로 자연 누적)
            with self._gap_lock:
                self._gap_inflight.discard(code)
            logger.warning("⚠️ 갭보정 큐 full → %s 스킵 (WS 실시간 누적으로 대체)", code)

    def _trigger_bulk_refill_async(self) -> None:
        """WS 재접속 시 현재 구독된 전 종목의 갭보정 완료 마커를 리셋하고 재큐잉."""
        if not (self.ws_cache and self.candle_agg):
            return
        if self._bulk_refill_running:
            return
        self._bulk_refill_running = True

        def _bulk():
            try:
                with self.ws_cache._sub_lock:
                    codes = sorted(self.ws_cache._subscribed)
                # 재접속이므로 모든 종목 갭보정 재실행
                with self._gap_lock:
                    self._gap_filled.clear()
                logger.info(
                    "🔄 [갭보정-전체] WS 재접속 → %d종목 큐 재등록", len(codes),
                )
                for code in codes:
                    self._enqueue_gap_fill(code)
            finally:
                self._bulk_refill_running = False

        threading.Thread(target=_bulk, name="WS-BulkRefill", daemon=True).start()

    def _gap_worker_loop(self) -> None:
        """단일 워커 루프: 큐에서 코드 꺼내 순차 처리 → REST 레이트리밋 자연 완충."""
        # 크레덴셜은 첫 작업 시점에 1회 조회 후 캐시 (env 변경 무시하고 세션 유지)
        kw_key = kw_secret = None
        kw_mock = False
        kw_resolved = False

        while True:
            try:
                code = self._gap_q.get(timeout=1.0)
            except queue.Empty:
                continue

            if code is None:  # 종료 시그널
                return

            # 장중만 실행 (장외면 완료 마커 찍고 다음)
            if not self._is_market_hours() and not get_env_bool("WS_GAP_FILL_OFF_HOURS", False):
                with self._gap_lock:
                    self._gap_inflight.discard(code)
                    self._gap_filled.add(code)
                self._gap_q.task_done()
                continue

            if not kw_resolved:
                kw_key, kw_secret, kw_mock = self._get_kiwoom_credentials()
                use_kiwoom = bool(kw_key and kw_secret and get_kiwoom_candles_df is not None)
                if use_kiwoom:
                    kw_status = f"✅ ({'모의' if kw_mock else '실전'})"
                else:
                    kw_status = "❌"
                logger.info(
                    "🔧 [갭보정-워커] kiwoom=%s, KIS_fallback=%s",
                    kw_status,
                    "ON" if get_env_bool("WS_GAP_FILL_KIS_FALLBACK", False) else "OFF",
                )
                kw_resolved = True

            try:
                self._fill_gap_for_code(
                    code,
                    kw_key=kw_key, kw_secret=kw_secret, kw_mock=kw_mock,
                )
            except Exception as e:
                logger.debug("갭보정 워커 예외 (%s): %s", code, e)
            finally:
                with self._gap_lock:
                    self._gap_inflight.discard(code)
                    self._gap_filled.add(code)
                self._gap_q.task_done()

            # 종목 간 짧은 sleep (REST 레이트리밋 완충)
            time.sleep(random.uniform(0.2, 0.4))

    @staticmethod
    def _is_market_hours() -> bool:
        """
        09:00~15:30 KST 평일 여부.
        KIS 모의투자 서버는 장외시간 inquire-time-itemchartprice 호출에
        HTTP 500 을 던지므로, 갭보정 REST 호출은 장중에만 시도한다.
        (키움 ka10080 은 장외에도 동작하지만, 매매 자체가 장중에만 의미 있으므로 통일)
        """
        now = time.localtime()
        if now.tm_wday >= 5:  # 토/일
            return False
        hhmm = now.tm_hour * 100 + now.tm_min
        return 900 <= hhmm <= 1530

    def _get_kiwoom_credentials(self):
        """
        키움 분봉 갭보정용 키 조회 — **KIS 처럼 모의/실전 토글 가능**.

        토글 결정 우선순위
        ------------------
          1) ``KIWOOM_MOCK`` (있으면 단독 사용 — 키움만 별도 토글하고 싶을 때)
          2) 미지정 시 ``KIS_MOCK`` 폴백 (한 번만 설정해도 동기화)

        키 슬롯 매핑
        -----------
          mock=True  → ``KIWOOM_APP_KEY_MOCK``  → 없으면 ``KIWOOM_APP_KEY`` (레거시) 폴백
          mock=False → ``KIWOOM_APP_KEY_REAL``  → 없으면 ``KIWOOM_APP_KEY`` (레거시) 폴백

        키 없으면 ``(None, None, is_mock)`` 반환 → 키움 비활성, KIS REST fallback 사용
        (``WS_GAP_FILL_KIS_FALLBACK`` 권장 ON)

        주의
        ----
          레거시 폴백은 키-도메인이 어긋나면 키움이 ``8030`` 으로 거부한다.
          (예: 모의 전용 키를 실전 도메인에 던지면 8030.) 폴백 사용 시 로그
          한 줄로 명시한다.

        Returns:
            (app_key, app_secret, is_mock)
        """
        if get_kiwoom_candles_df is None:
            return None, None, False
        try:
            # ── 1. 토글 결정 ──────────────────────────────────────────
            kw_mock_raw = (get_env_from_db("KIWOOM_MOCK", "") or "").strip().lower()
            if kw_mock_raw in ("true", "1", "yes", "y", "on"):
                is_mock = True
            elif kw_mock_raw in ("false", "0", "no", "n", "off"):
                is_mock = False
            else:
                is_mock = get_env_bool("KIS_MOCK", True)

            # ── 2. 키 슬롯 선택 (모의/실전) ───────────────────────────
            if is_mock:
                kw_key = (get_env_from_db("KIWOOM_APP_KEY_MOCK", "") or "").strip()
                kw_secret = (get_env_from_db("KIWOOM_APP_SECRET_MOCK", "") or "").strip()
            else:
                kw_key = (get_env_from_db("KIWOOM_APP_KEY_REAL", "") or "").strip()
                kw_secret = (get_env_from_db("KIWOOM_APP_SECRET_REAL", "") or "").strip()

            # ── 3. 레거시 단일 필드 폴백 (KIWOOM_APP_KEY/_SECRET) ────
            if not kw_key or not kw_secret:
                legacy_key = (get_env_from_db("KIWOOM_APP_KEY", "") or "").strip()
                legacy_secret = (get_env_from_db("KIWOOM_APP_SECRET", "") or "").strip()
                if legacy_key and legacy_secret:
                    kw_key = kw_key or legacy_key
                    kw_secret = kw_secret or legacy_secret
                    logger.info(
                        "🔧 [키움] %s 슬롯 비어있어 LEGACY KIWOOM_APP_KEY 폴백 사용 "
                        "(키-도메인 불일치 시 8030 발생 가능)",
                        "MOCK" if is_mock else "REAL",
                    )

            if not kw_key or not kw_secret:
                return None, None, is_mock

            return kw_key, kw_secret, is_mock
        except Exception as e:
            logger.debug("키움 크레덴셜 조회 예외: %s", e)
            return None, None, False

    def _fill_gap_for_code(
        self,
        code: str,
        *,
        kw_key: Optional[str] = None,
        kw_secret: Optional[str] = None,
        kw_mock: bool = False,
    ) -> None:
        """
        단일 종목 갭 보정 — 워커 스레드 전용 (전략 쓰레드에서 직접 호출 금지).

        정책 (기존 kis_scalping_ver2._fill_all_gaps 개선):
          [1] 키움 ka10080 우선 (1/3/15/60분봉 native + 과거봉 확보)
          [2] 키움 실패 시 KIS fallback — **기본 OFF** (`WS_GAP_FILL_KIS_FALLBACK=false`)
              → KIS 모의 서버가 장중에도 HTTP 500 을 자주 반환해 로그 오염 + 백오프 지연.
                 키움 있으면 굳이 안 쳐도 됨. WS 틱이 쌓여 자연 보완됨.
              → env 로 true 지정 시에만 tf<=3 한정 KIS 호출.
          [3] 어느 경로든 실패 → WS 실시간 틱으로 자연 누적 (CandleAggregator)
        """
        if not (self.ws_cache and self.candle_agg):
            return
        if not self._is_market_hours() and not get_env_bool("WS_GAP_FILL_OFF_HOURS", False):
            return

        limit = get_env_int("WS_GAP_FILL_LIMIT", 120)
        use_kiwoom = bool(kw_key and kw_secret and get_kiwoom_candles_df is not None)
        kis_fallback_on = get_env_bool("WS_GAP_FILL_KIS_FALLBACK", False)

        for tf in self.candle_agg.timeframes:
            df = None

            if use_kiwoom:
                try:
                    df = get_kiwoom_candles_df(
                        code, tf, kw_key, kw_secret,
                        is_mock=kw_mock, n=limit,
                    )
                except Exception as e:
                    logger.debug("키움 갭보정 실패 (%s %dM): %s", code, tf, e)

            # KIS fallback — env 로 명시적 ON 일 때만 (1/3분봉 한정)
            if (df is None or df.empty) and kis_fallback_on and tf <= 3:
                try:
                    df = self.kis_client.get_minute_chart(
                        code, period=str(tf), limit=limit,
                    )
                except Exception as e:
                    logger.debug("KIS 갭보정 실패 (%s %dM): %s", code, tf, e)

            if df is not None and not df.empty:
                self.candle_agg.fill_gap_from_rest(code, tf, df)

            # 같은 종목 내 타임프레임 전환 사이 짧은 sleep (차트 API 레이트리밋)
            time.sleep(random.uniform(0.15, 0.3))

    # ------------------------------------------------------------------
    # 영구 구독 / 타임프레임 해석
    # ------------------------------------------------------------------
    def _load_permanent_codes(self) -> None:
        raw = get_env_from_db("PERMANENT_WS_CODES", "069500,229200")
        self._permanent_codes = {c.strip() for c in str(raw).split(",") if c.strip()}

    def _resolve_timeframes(self) -> List[int]:
        """SCALP 1분 + SHORT 3분 + 추세 15/60분. env 로 확장 가능."""
        tf_raw = get_env_from_db("WS_TIMEFRAMES", "1,3,15,60")
        try:
            tfs = [int(x.strip()) for x in str(tf_raw).split(",") if x.strip()]
            # 최소 1분, 3분은 포함 보장 (두 전략 필수 TF)
            for must in (1, 3):
                if must not in tfs:
                    tfs.append(must)
            return sorted(set(tfs))
        except Exception:
            return [1, 3, 15, 60]
