kis_trader/scripts/query_kiwoom_condition_snapshot.py

#!/usr/bin/env python3
"""
키움 저장조건식 → 현재 매칭 종목 1회 조회 (HTS 실검 대조용)
============================================================

실매 봇 RAM/히스토리 sticky 가 아니라, 키움 WS 로 CNSRREQ 초기 응답을
받아 **지금 서버 조건식이 뽑는 종목**을 출력한다.

흐름:
  LOGIN → CNSRLST → (이름/seq로) CNSRREQ → 초기 data 출력 → CNSRCLR → 종료

주의:
  - 실매 ``kis_trader_main`` 과 **같은 키움 계정** 을 쓴다.
    이미 실시간 등록된 seq 면 900003 이 날 수 있어, 1회 CLR→REQ 후 다시 CLR 한다.
  - 장후에는 편입/이탈 push 가 거의 없고, 초기 리스트만 의미 있다.
  - 주문·조건식 저장은 하지 않는다.

사용:
  cd /home/hoon/kis_bot
  python3 -u scripts/query_kiwoom_condition_snapshot.py
  python3 -u scripts/query_kiwoom_condition_snapshot.py --name momentum
  python3 -u scripts/query_kiwoom_condition_snapshot.py --seq 3
  python3 -u scripts/query_kiwoom_condition_snapshot.py --name momentum --compare-db

로그:
  logs/kiwoom_cond_snapshot_YYYYMMDD_HHMMSS.log
"""
from __future__ import annotations

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

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

from kis_trader.utils.env import get_env_from_db  # noqa: E402
from kis_trader.ws.kis_ws import _get_kiwoom_token_cached  # noqa: E402


def _load_kiwoom_creds() -> Tuple[str, str, bool]:
    kw_mock_raw = (get_env_from_db("KIWOOM_MOCK", "") or "").strip().lower()
    is_mock = kw_mock_raw in ("1", "true", "y", "yes", "on")
    if is_mock:
        key = (get_env_from_db("KIWOOM_APP_KEY_MOCK", "") or "").strip()
        secret = (get_env_from_db("KIWOOM_APP_SECRET_MOCK", "") or "").strip()
    else:
        key = (get_env_from_db("KIWOOM_APP_KEY_REAL", "") or "").strip()
        secret = (get_env_from_db("KIWOOM_APP_SECRET_REAL", "") or "").strip()
    if not key or not secret:
        key = key or (get_env_from_db("KIWOOM_APP_KEY", "") or "").strip()
        secret = secret or (get_env_from_db("KIWOOM_APP_SECRET", "") or "").strip()
    return key, secret, is_mock


def _seq_name(item: Any) -> Tuple[str, str]:
    if isinstance(item, (list, tuple)):
        return (
            str(item[0]) if len(item) > 0 else "",
            str(item[1]) if len(item) > 1 else "",
        )
    if isinstance(item, dict):
        return str(item.get("seq") or ""), str(item.get("name") or "")
    return "", ""


def _extract_code(item: Any) -> str:
    raw = ""
    if isinstance(item, dict):
        raw = str(item.get("9001") or item.get("jmcode") or item.get("code") or "").strip()
    elif isinstance(item, (list, tuple)) and item:
        raw = str(item[0]).strip()
    if not raw:
        return ""
    # 키움 조건검색 응답은 종종 'A005930' 형태 — 비교·표시용으로 A 접두 제거
    if len(raw) >= 7 and raw[0] in ("A", "a") and raw[1:].isdigit():
        return raw[1:]
    if raw.upper().startswith("A") and len(raw) == 7:
        return raw[1:]
    return raw


def _resolve_names(codes: List[str]) -> Dict[str, str]:
    """DB 히스토리에서 종목명 보강 (없으면 코드 그대로)."""
    out: Dict[str, str] = {c: c for c in codes}
    if not codes:
        return out
    try:
        from database import TradeDB

        db = TradeDB()
        try:
            for c in codes:
                row = db.conn.execute(
                    "SELECT name FROM target_candidates_history "
                    "WHERE code=%s AND name IS NOT NULL AND name<>'' AND name<>code "
                    "ORDER BY id DESC LIMIT 1",
                    (c,),
                ).fetchone()
                if row and row.get("name"):
                    out[c] = str(row["name"])
                    continue
                row2 = db.conn.execute(
                    "SELECT name FROM ls_candidates_history "
                    "WHERE code=%s AND name IS NOT NULL AND name<>'' AND name<>code "
                    "ORDER BY id DESC LIMIT 1",
                    (c,),
                ).fetchone()
                if row2 and row2.get("name"):
                    out[c] = str(row2["name"])
        finally:
            db.close()
    except Exception as e:
        print(f"⚠️ 종목명 DB 보강 스킵: {e}")
    return out


def _latest_db_universe(strategy_id: str = "MOMENTUM") -> Tuple[str, List[str]]:
    from database import TradeDB

    db = TradeDB()
    try:
        slot = db.conn.execute(
            "SELECT slot_key FROM target_candidates_history "
            "WHERE strategy_id=%s GROUP BY slot_key ORDER BY slot_key DESC LIMIT 1",
            (strategy_id,),
        ).fetchone()
        if not slot:
            return "", []
        sk = str(slot["slot_key"])
        rows = db.conn.execute(
            "SELECT code FROM target_candidates_history "
            "WHERE strategy_id=%s AND slot_key=%s ORDER BY code",
            (strategy_id, sk),
        ).fetchall()
        return sk, [str(r["code"]) for r in rows]
    finally:
        db.close()


def query_snapshot(
    *,
    name: str = "momentum",
    seq: str = "",
    search_type: str = "1",
    timeout_sec: float = 25.0,
    log_path: Optional[Path] = None,
) -> Dict[str, Any]:
    try:
        import websocket
    except Exception as e:
        raise RuntimeError(f"websocket-client 미설치: {e}") from e

    key, secret, is_mock = _load_kiwoom_creds()
    if not key or not secret:
        raise RuntimeError("키움 앱키/시크릿 미설정 (KIWOOM_APP_KEY_REAL 등)")

    token = _get_kiwoom_token_cached(key, secret, is_mock)
    if not token:
        raise RuntimeError("키움 토큰 발급 실패 (au10001)")

    url = (
        get_env_from_db("KIWOOM_WS_URL_MOCK", "wss://mockapi.kiwoom.com:10000/api/dostk/websocket")
        if is_mock
        else get_env_from_db("KIWOOM_WS_URL_REAL", "wss://api.kiwoom.com:10000/api/dostk/websocket")
    )

    st: Dict[str, Any] = {
        "seq": (seq or "").strip(),
        "name": (name or "").strip(),
        "conditions": [],
        "codes": [],
        "error": None,
        "done": False,
        "cleared": False,
        "retried_900003": False,
        "raw_cnsrreq": None,
        "mock": is_mock,
        "url": url,
    }
    lock = threading.Lock()

    def _log(msg: str) -> None:
        line = f"[{datetime.now().strftime('%H:%M:%S')}] {msg}"
        print(line, flush=True)
        if log_path:
            with open(log_path, "a", encoding="utf-8") as f:
                f.write(line + "\n")

    def _send(ws, payload: dict) -> None:
        ws.send(json.dumps(payload))

    def on_open(ws):
        _log("LOGIN 발송")
        _send(ws, {"trnm": "LOGIN", "token": token})

    def on_message(ws, message):
        try:
            data = json.loads(message)
        except Exception:
            _log(f"non-json: {str(message)[:200]}")
            return
        trnm = data.get("trnm")
        if trnm == "PING":
            try:
                ws.send(message)
            except Exception:
                pass
            return

        if trnm == "LOGIN":
            rc = str(data.get("return_code"))
            if rc not in ("0", "0.0"):
                st["error"] = f"LOGIN 실패 rc={rc} msg={data.get('return_msg')}"
                ws.close()
                return
            _log("LOGIN OK → CNSRLST")
            _send(ws, {"trnm": "CNSRLST"})
            return

        if trnm == "CNSRLST":
            items = data.get("data") or []
            st["conditions"] = [
                {"seq": a, "name": b} for a, b in (_seq_name(it) for it in items) if a or b
            ]
            _log(
                "저장조건식 %d개: %s"
                % (
                    len(st["conditions"]),
                    ", ".join(f"{c['seq']}:{c['name']}" for c in st["conditions"]),
                )
            )
            target_seq = st["seq"]
            if not target_seq:
                want = st["name"].lower()
                for c in st["conditions"]:
                    if str(c["name"]).strip().lower() == want:
                        target_seq = str(c["seq"])
                        break
            if not target_seq:
                st["error"] = f"조건식 없음 name={st['name']!r} seq={st['seq']!r}"
                ws.close()
                return
            st["seq"] = target_seq
            _log(
                f"CNSRREQ seq={target_seq} name={st['name'] or '?'} "
                f"search_type={search_type}"
            )
            _send(
                ws,
                {
                    "trnm": "CNSRREQ",
                    "seq": target_seq,
                    "search_type": str(search_type),
                    "stex_tp": "K",
                },
            )
            return

        if trnm == "CNSRREQ":
            st["raw_cnsrreq"] = data
            rc = str(data.get("return_code"))
            msg = str(data.get("return_msg") or "")
            if rc == "900003" and not st["retried_900003"]:
                st["retried_900003"] = True
                _log(f"900003 이미등록 → CNSRCLR 1회 후 재요청 seq={st['seq']}")
                _send(ws, {"trnm": "CNSRCLR", "seq": st["seq"]})
                return
            if rc not in ("0", "0.0"):
                st["error"] = f"CNSRREQ 실패 rc={rc} msg={msg}"
                try:
                    _send(ws, {"trnm": "CNSRCLR", "seq": st["seq"]})
                except Exception:
                    pass
                ws.close()
                return
            codes = [_extract_code(it) for it in (data.get("data") or [])]
            codes = sorted({c for c in codes if c})
            st["codes"] = codes
            _log(f"CNSRREQ OK — 매칭 {len(codes)}종목")
            _send(ws, {"trnm": "CNSRCLR", "seq": st["seq"]})
            return

        if trnm == "CNSRCLR":
            rc = str(data.get("return_code"))
            _log(f"CNSRCLR rc={rc}")
            # 900003 재시도 경로: CLR 후 다시 REQ
            if st["retried_900003"] and not st["codes"] and not st.get("_req_after_clr"):
                st["_req_after_clr"] = True
                time.sleep(0.8)
                _log(f"CLR 후 CNSRREQ 재발송 seq={st['seq']}")
                _send(
                    ws,
                    {
                        "trnm": "CNSRREQ",
                        "seq": st["seq"],
                        "search_type": str(search_type),
                        "stex_tp": "K",
                    },
                )
                return
            st["cleared"] = True
            st["done"] = True
            ws.close()
            return

        _log(f"기타 {trnm}: {json.dumps(data, ensure_ascii=False)[:240]}")

    def on_error(ws, err):
        with lock:
            if not st["error"]:
                st["error"] = f"WS 오류: {err}"
        _log(f"WS error: {err}")

    def on_close(ws, code, msg):
        st["done"] = True

    ws = websocket.WebSocketApp(
        url,
        on_open=on_open,
        on_message=on_message,
        on_error=on_error,
        on_close=on_close,
    )
    th = threading.Thread(
        target=ws.run_forever,
        kwargs={"ping_interval": 0},
        daemon=True,
    )
    th.start()

    deadline = time.time() + float(timeout_sec)
    while time.time() < deadline and not st["done"]:
        if st["error"] and not st.get("_req_after_clr"):
            # 에러 후 CLR 대기 중일 수 있음
            if "CNSRREQ 실패" in str(st["error"]):
                time.sleep(0.5)
                break
        time.sleep(0.15)

    try:
        ws.close()
    except Exception:
        pass
    th.join(timeout=3.0)

    if st["error"] and not st["codes"]:
        raise RuntimeError(st["error"])
    return st


def main() -> int:
    ap = argparse.ArgumentParser(description="키움 조건식 현재 매칭 종목 1회 조회")
    ap.add_argument("--name", default="momentum", help="조건식 이름 (기본 momentum)")
    ap.add_argument("--seq", default="", help="seq 직접 지정 시 이름 무시")
    ap.add_argument(
        "--search-type",
        default="1",
        help="CNSRREQ search_type (봇과 동일 기본 1)",
    )
    ap.add_argument("--timeout", type=float, default=25.0)
    ap.add_argument(
        "--compare-db",
        action="store_true",
        help="target_candidates_history 최신 MOMENTUM 슬롯과 비교",
    )
    ap.add_argument(
        "--strategy-id",
        default="MOMENTUM",
        help="--compare-db 시 history strategy_id",
    )
    args = ap.parse_args()

    log_dir = ROOT / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_path = log_dir / f"kiwoom_cond_snapshot_{ts}.log"

    print(f"📝 log: {log_path}")
    print(
        "⚠️ 실매 봇과 동일 키움 계정 — 조회 후 CNSRCLR 함. "
        "장중엔 봇 조건등록과 순간 충돌 가능."
    )

    try:
        st = query_snapshot(
            name=args.name,
            seq=args.seq,
            search_type=args.search_type,
            timeout_sec=args.timeout,
            log_path=log_path,
        )
    except Exception as e:
        print(f"❌ {e}")
        return 1

    codes: List[str] = list(st.get("codes") or [])
    names = _resolve_names(codes)

    print()
    print("=" * 60)
    print(
        f"키움 조건식 스냅샷  name={args.name!r} seq={st.get('seq')}  "
        f"매칭={len(codes)}종  mock={st.get('mock')}"
    )
    print("=" * 60)
    if not codes:
        print("(매칭 종목 없음)")
    else:
        for i, c in enumerate(codes, 1):
            print(f"  {i:2d}. {c}  {names.get(c, c)}")
    print("=" * 60)

    if args.compare_db:
        sk, db_codes = _latest_db_universe(args.strategy_id)
        db_set: Set[str] = set(db_codes)
        live_set: Set[str] = set(codes)
        only_api = sorted(live_set - db_set)
        only_db = sorted(db_set - live_set)
        both = sorted(live_set & db_set)
        print()
        print(f"[DB 비교] strategy={args.strategy_id} latest_slot={sk} n={len(db_codes)}")
        print(f"  교집합 {len(both)}: {both}")
        print(f"  API만 {len(only_api)}: {only_api}")
        print(f"  DB만(sticky) {len(only_db)}: {only_db}")

    with open(log_path, "a", encoding="utf-8") as f:
        f.write(
            json.dumps(
                {
                    "seq": st.get("seq"),
                    "name": args.name,
                    "codes": codes,
                    "names": names,
                    "conditions": st.get("conditions"),
                },
                ensure_ascii=False,
                indent=2,
            )
            + "\n"
        )
    print(f"✅ 완료 — {log_path}")
    return 0


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