#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
KIS WebSocket(H0STCNT0) 연결·구독 진단 스크립트
================================================
봇 로그의 "장외 서버 거부 추정" / 즉시 끊김이 **진짜 KIS 차단**인지,
**구독 폭주·중복 세션·구독 오류**인지 단계별로 검증한다.

⚠️ kis_trader_main.service 가 이미 KIS WS 를 쓰는 중이면
   동일 appkey 로 2번째 WS 를 열 때 기존 세션이 끊길 수 있다.
   --skip-if-service-running (기본) 으로 서비스 기동 중이면 Phase 2+ 를 스킵한다.

사용:
  cd /home/hoon/kis_bot
  python3 kis_trader/scripts/test_kis_ws_diagnostic.py
  python3 kis_trader/scripts/test_kis_ws_diagnostic.py --force --codes 005930,069500
  python3 kis_trader/scripts/test_kis_ws_diagnostic.py --phase duplicate_sub

로그 파일:
  logs/test_kis_ws_diagnostic.log  (동시에 stdout 출력)
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import subprocess
import sys
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

HERE = Path(__file__).resolve()
ROOT = HERE.parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import requests  # noqa: E402

from kis_trader.execution.kis_client import KISClient  # noqa: E402
from kis_trader.utils.env import get_env_bool, get_env_from_db  # noqa: E402

LOG_DIR = ROOT / "logs"
LOG_DIR.mkdir(exist_ok=True)
LOG_PATH = LOG_DIR / "test_kis_ws_diagnostic.log"

DEFAULT_CODES = ["000660", "005380", "005930", "069500", "229200", "379810", "466930"]


def _setup_logging() -> logging.Logger:
    lg = logging.getLogger("test_kis_ws")
    lg.setLevel(logging.DEBUG)
    lg.handlers.clear()
    fmt = logging.Formatter("[%(asctime)s] %(message)s", datefmt="%H:%M:%S")
    for h in (logging.StreamHandler(sys.stdout), logging.FileHandler(LOG_PATH, encoding="utf-8")):
        h.setFormatter(fmt)
        lg.addHandler(h)
    return lg


def _is_market_hours(is_mock: bool) -> bool:
    now = datetime.now()
    if now.weekday() >= 5:
        return False
    open_h, open_m = (9, 0) if is_mock else (8, 25)
    from datetime import time as dtime
    return dtime(open_h, open_m) <= now.time() <= dtime(16, 5)


def _service_running(name: str = "kis_trader_main.service") -> bool:
    try:
        r = subprocess.run(
            ["systemctl", "is-active", name],
            capture_output=True,
            text=True,
            timeout=5,
        )
        return r.stdout.strip() == "active"
    except Exception:
        return False


def _fetch_approval_key(base_url: str, app_key: str, app_secret: str) -> Tuple[Optional[str], Dict[str, Any]]:
    url = f"{base_url}/oauth2/Approval"
    body = {
        "grant_type": "client_credentials",
        "appkey": app_key,
        "secretkey": app_secret,
    }
    try:
        r = requests.post(url, json=body, timeout=10)
        data = r.json()
        return data.get("approval_key"), data
    except Exception as e:
        return None, {"error": str(e)}


class WsProbe:
    """단일 WS 연결 프로브 — 수신 메시지·끊김 시각 기록."""

    def __init__(
        self,
        ws_url: str,
        approval_key: str,
        logger: logging.Logger,
        gap_sec: float = 0.0,
    ):
        self.ws_url = ws_url
        self.approval_key = approval_key
        self.logger = logger
        self.gap_sec = gap_sec
        self.messages: List[str] = []
        self.json_responses: List[Dict[str, Any]] = []
        self.ticks = 0
        self.connected_at: float = 0.0
        self.closed_at: float = 0.0
        self.close_code: Any = None
        self.close_msg: str = ""
        self.error: str = ""
        self._ws = None
        self._thread: Optional[threading.Thread] = None
        self._done = threading.Event()
        self._sent_codes: List[str] = []

    def _build_sub(self, code: str, subscribe: bool = True) -> str:
        return json.dumps({
            "header": {
                "approval_key": self.approval_key,
                "custtype": "P",
                "tr_type": "1" if subscribe else "2",
                "content-type": "utf-8",
            },
            "body": {
                "input": {
                    "tr_id": "H0STCNT0",
                    "tr_key": code,
                }
            },
        })

    def _on_open(self, ws) -> None:
        self.connected_at = time.time()
        self.logger.info("  [WS] on_open OK (t=0)")

    def _on_message(self, ws, message: str) -> None:
        self.messages.append(message)
        raw = (message or "").strip()
        if raw == "PINGPONG":
            try:
                ws.send("PINGPONG")
            except Exception:
                pass
            self.logger.debug("  [WS] PINGPONG echo")
            return
        if raw.startswith("{"):
            try:
                j = json.loads(raw)
                self.json_responses.append(j)
                hdr = j.get("header") or {}
                body = j.get("body") or {}
                rt = str(body.get("rt_cd", ""))
                msg1 = str(body.get("msg1", ""))
                tr_id = hdr.get("tr_id", "")
                self.logger.info(
                    "  [WS] JSON tr_id=%s rt_cd=%s msg1=%s",
                    tr_id, rt, msg1,
                )
                if rt and rt != "0":
                    self.logger.warning("  [WS] ★ 구독/서버 오류 rt_cd=%s: %s", rt, msg1)
            except Exception as e:
                self.logger.warning("  [WS] JSON 파싱 실패: %s | raw=%s", e, raw[:200])
            return
        parts = raw.split("|")
        if len(parts) >= 2 and parts[1] == "H0STCNT0":
            self.ticks += 1
            if self.ticks <= 3:
                self.logger.info("  [WS] H0STCNT0 tick #%d (샘플)", self.ticks)

    def _on_error(self, ws, error) -> None:
        self.error = str(error)
        self.logger.warning("  [WS] on_error: %s", error)

    def _on_close(self, ws, code, msg) -> None:
        self.closed_at = time.time()
        self.close_code = code
        self.close_msg = str(msg or "")
        dur = self.closed_at - self.connected_at if self.connected_at else 0.0
        self.logger.info(
            "  [WS] on_close code=%s dur=%.2fs msg=%s",
            code, dur, self.close_msg or "-",
        )
        self._done.set()

    def run(
        self,
        codes: List[str],
        *,
        duplicate_subscribe: bool = False,
        hold_sec: float = 8.0,
    ) -> Dict[str, Any]:
        try:
            import websocket as ws_lib
        except ImportError:
            return {"ok": False, "reason": "websocket-client 미설치"}

        def _runner():
            app = ws_lib.WebSocketApp(
                self.ws_url,
                on_open=self._on_open,
                on_message=self._on_message,
                on_error=self._on_error,
                on_close=self._on_close,
            )
            self._ws = app

            def _subscribe_all():
                time.sleep(0.3)
                for i, code in enumerate(codes):
                    if self.gap_sec > 0 and i > 0:
                        time.sleep(self.gap_sec)
                    try:
                        app.send(self._build_sub(code, True))
                        self._sent_codes.append(code)
                        self.logger.info("  [WS] subscribe sent: %s", code)
                    except Exception as e:
                        self.logger.warning("  [WS] subscribe fail %s: %s", code, e)
                if duplicate_subscribe and codes:
                    self.logger.info("  [WS] duplicate subscribe burst (봇 _on_open 버그 재현)")
                    for code in codes:
                        try:
                            app.send(self._build_sub(code, True))
                        except Exception:
                            pass

            if codes:
                threading.Thread(target=_subscribe_all, daemon=True).start()

            app.run_forever(ping_interval=20, ping_timeout=10)

        self._thread = threading.Thread(target=_runner, daemon=True)
        self._thread.start()
        self._done.wait(timeout=hold_sec + 15.0)
        gap = max(0.0, float(self.gap_sec or 0.12))
        if self._ws and self._sent_codes:
            self.logger.info("  [WS] unsubscribe %d 후 close", len(self._sent_codes))
            for i, code in enumerate(self._sent_codes):
                if i > 0 and gap > 0:
                    time.sleep(gap)
                try:
                    self._ws.send(self._build_sub(code, False))
                except Exception:
                    pass
            if gap > 0:
                time.sleep(gap)
        try:
            if self._ws:
                self._ws.close()
        except Exception:
            pass
        if self._thread.is_alive():
            self._thread.join(timeout=3)

        dur = (self.closed_at or time.time()) - self.connected_at if self.connected_at else 0.0
        instant = dur > 0 and dur < 3.0
        err_json = [j for j in self.json_responses if str((j.get("body") or {}).get("rt_cd", "0")) != "0"]
        return {
            "ok": self.connected_at > 0 and not instant and self.error == "",
            "connected": self.connected_at > 0,
            "duration_sec": round(dur, 2),
            "instant_drop": instant,
            "error": self.error,
            "close_code": self.close_code,
            "ticks": self.ticks,
            "json_errors": len(err_json),
            "json_total": len(self.json_responses),
            "messages": len(self.messages),
        }


def _verdict(phase: str, result: Dict[str, Any], market_open: bool) -> str:
    if result.get("reason"):
        return f"FAIL — {result['reason']}"
    if not result.get("connected"):
        return "FAIL — TCP/WS 핸드셰이크 실패 (approval_key·URL·방화벽 확인)"
    if result.get("json_errors", 0) > 0:
        return "FAIL — KIS JSON rt_cd≠0 (구독 거부·잘못된 tr_key·한도 초과)"
    if result.get("instant_drop"):
        if not market_open:
            return "EXPECTED — 장외 즉시 끊김 (KIS WS 서비스 시간 외)"
        if phase == "duplicate_sub":
            return "LIKELY — 중복 구독 폭주 후 즉시 끊김 (봇 _on_open 이중 전송 의심)"
        return "FAIL — 장중인데 즉시 끊김 (동시 세션·서버 장애·rate limit 의심)"
    if result.get("ticks", 0) > 0:
        return "OK — 연결 유지 + H0STCNT0 틱 수신"
    if result.get("duration_sec", 0) >= 5:
        return "OK — 연결 유지 (장중 틱 없음=거래 없는 종목·장외 가능)"
    return "WARN — 연결 짧음, 추가 확인 필요"


def main() -> int:
    parser = argparse.ArgumentParser(description="KIS WebSocket 진단")
    parser.add_argument(
        "--codes",
        default=",".join(DEFAULT_CODES),
        help="구독 테스트 종목 (쉼표구분)",
    )
    parser.add_argument(
        "--phase",
        choices=("all", "approval", "connect_only", "single", "multi", "duplicate_sub"),
        default="all",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="kis_trader_main.service 기동 중에도 WS 테스트 실행",
    )
    parser.add_argument(
        "--gap-sec",
        type=float,
        default=0.0,
        help="종목별 구독 간격(초). 0=봇과 동일 즉시 연속",
    )
    parser.add_argument(
        "--hold-sec",
        type=float,
        default=8.0,
        help="연결 유지 관측 시간(초)",
    )
    args = parser.parse_args()
    codes = [c.strip() for c in args.codes.split(",") if c.strip()]

    log = _setup_logging()
    log.info("=" * 60)
    log.info("KIS WebSocket 진단 시작 → log: %s", LOG_PATH)
    log.info("=" * 60)

    mock = get_env_bool("KIS_MOCK", True)
    client = KISClient(mock=mock)
    ws_url = (
        get_env_from_db("KIS_WS_URL_MOCK", "ws://ops.koreainvestment.com:31000")
        if mock
        else get_env_from_db("KIS_WS_URL_REAL", "ws://ops.koreainvestment.com:21000")
    )
    market_open = _is_market_hours(mock)
    svc = _service_running()

    log.info("KIS_MOCK=%s | WS_URL=%s", mock, ws_url)
    log.info("장 서비스 시간=%s (mock기준 open %s)", market_open, "09:00" if mock else "08:25")
    log.info("kis_trader_main.service active=%s", svc)
    if svc and not args.force:
        log.warning(
            "⚠️ 메인 봇이 WS 사용 중 → 2번째 연결은 기존 세션을 끊을 수 있음. "
            "Phase connect/subscribe 는 스킵 (--force 로 강행)",
        )

    # Phase 1: approval_key
    if args.phase in ("all", "approval"):
        log.info("--- Phase 1: approval_key REST ---")
        key, raw = _fetch_approval_key(client.base_url, client.app_key, client.app_secret)
        if key:
            log.info("  approval_key OK (앞8자 %s…)", key[:8])
        else:
            log.error("  approval_key FAIL: %s", raw)
            return 1

    if svc and not args.force and args.phase != "approval":
        log.info("진단 종료 (서비스 기동 중, WS Phase 스킵). --force 로 재실행하세요.")
        return 0

    key, _ = _fetch_approval_key(client.base_url, client.app_key, client.app_secret)
    if not key:
        log.error("approval_key 없음 — 종료")
        return 1

    phases: List[Tuple[str, List[str], bool]] = []
    if args.phase in ("all", "connect_only"):
        phases.append(("connect_only", [], False))
    if args.phase in ("all", "single"):
        phases.append(("single", [codes[0] if codes else "005930"], False))
    if args.phase in ("all", "multi"):
        phases.append(("multi", codes, False))
    if args.phase in ("all", "duplicate_sub"):
        phases.append(("duplicate_sub", codes[:3] if codes else ["005930"], True))

    if args.phase == "duplicate_sub":
        phases = [("duplicate_sub", codes[:3] if codes else ["005930"], True)]

    summary: List[str] = []
    for name, sub_codes, dup in phases:
        log.info("--- Phase: %s (codes=%d dup=%s gap=%.1fs hold=%.0fs) ---",
                 name, len(sub_codes), dup, args.gap_sec, args.hold_sec)
        probe = WsProbe(ws_url, key, log, gap_sec=args.gap_sec)
        res = probe.run(sub_codes, duplicate_subscribe=dup, hold_sec=args.hold_sec)
        v = _verdict(name, res, market_open)
        log.info("  결과: %s | detail=%s", v, res)
        summary.append(f"{name}: {v}")

    log.info("=" * 60)
    log.info("요약")
    for line in summary:
        log.info("  • %s", line)
    log.info("")
    log.info("해석 가이드:")
    log.info("  • 장중 instant_drop + 서비스 active → 동시 WS 세션 충돌 가능성 큼")
    log.info("  • json_errors>0 → KIS 가 구독 거부 (rt_cd/msg1 로그 확인)")
    log.info("  • duplicate_sub 만 instant_drop → kis_ws._on_open 이중 subscribe 버그")
    log.info("  • 장외 instant_drop → '장외 서버 거부' 메시지는 이 경우 정상")
    log.info("=" * 60)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())