kis_trader/permanent_subs.py

"""
permanent_subs.py
=================
영구구독(Permanent WS Subscriptions) 통합 테이블 — **국내(KR)·해외(US) 단일 소스**.

배경:
- 기존 ``env_config.PERMANENT_WS_CODES`` 는 콤마 문자열이라 **시장/거래소/심볼 메타를
  담지 못함.** 해외(QQQM 등)는 국내 WS(H0STCNT0)로는 못 받고, 거래소 코드가 필요하다.
- 이 테이블은 코드별로 ``market_type/exchange/symbol/tf_min`` 을 보관해
  - KR → 국내 실시간 WS(H0STCNT0)
  - US → 해외 실시간 WS(HDFSCNT0, tr_key=D{EXCD}{SYMBOL})
  로 각각 라우팅할 수 있게 한다.

호환:
- ``env_config.PERMANENT_WS_CODES`` 는 폴백으로 유지(테이블이 비어 있을 때만).
- ``migrate_env_codes()`` 로 기존 콤마 코드를 1회 이관(중복은 건너뜀).
"""
from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

logger = logging.getLogger("permanent_subs")

# 해외(US) 실시간 WS tr_key 거래소 코드 — 주문용(NASD/NYSE/AMEX) → 시세용(NAS/NYS/AMS)
_US_EXCD_MAP = {
    "NASD": "NAS", "NAS": "NAS",
    "NYSE": "NYS", "NYS": "NYS",
    "AMEX": "AMS", "AMS": "AMS",
}

_PERM_DDL = """
CREATE TABLE IF NOT EXISTS permanent_subscriptions (
    id          BIGINT AUTO_INCREMENT PRIMARY KEY,
    code        VARCHAR(16)  NOT NULL,
    market_type VARCHAR(8)   NOT NULL DEFAULT 'KR',
    exchange    VARCHAR(16)  NOT NULL DEFAULT 'KRX',
    symbol      VARCHAR(32)  NOT NULL DEFAULT '',
    tf_min      INT          NOT NULL DEFAULT 1,
    enabled     TINYINT      NOT NULL DEFAULT 1,
    note        VARCHAR(100) NOT NULL DEFAULT '',
    created_at  DATETIME     DEFAULT CURRENT_TIMESTAMP,
    UNIQUE KEY uq_perm_code (code)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='영구구독(국내 WS + 해외 WS) 종목'
"""


def _core(db: Any):
    """TradeDBExt 이면 ``.raw``, 아니면 그대로."""
    return getattr(db, "raw", db)


def ensure_permanent_subs_table(db: Any) -> None:
    raw = _core(db)
    raw.conn.execute(_PERM_DDL.strip())
    raw.conn.commit()


def classify_market(code: str) -> tuple[str, str, str]:
    """
    코드 형태로 (market_type, exchange, symbol) 추정.
    - 6자리 숫자 → KR/KRX
    - 영문 1~8자 → US/NASD (거래소는 이후 UI/DB에서 보정 가능)
    """
    c = str(code or "").strip().upper()
    if c.isdigit() and len(c) == 6:
        return "KR", "KRX", c
    if c.isalpha() and 1 <= len(c) <= 8:
        return "US", "NASD", c
    # 알 수 없으면 KR 취급(보수적)
    return "KR", "KRX", c


def us_ws_tr_key(exchange: str, symbol: str) -> str:
    """해외 실시간 WS tr_key — D + 시세거래소(NAS/NYS/AMS) + 심볼 (예: DNASQQQM)."""
    ex = str(exchange or "NASD").strip().upper()
    excd = _US_EXCD_MAP.get(ex, "NAS")
    sym = str(symbol or "").strip().upper()
    return f"D{excd}{sym}"


def upsert_permanent_sub(
    db: Any,
    code: str,
    market_type: Optional[str] = None,
    exchange: Optional[str] = None,
    symbol: Optional[str] = None,
    tf_min: int = 1,
    enabled: bool = True,
    note: str = "",
) -> None:
    """코드 1건 등록/갱신 (code UNIQUE).

    tf_min 은 저장·표시용 기준봉. WS는 틱→1분 집계가 기본이고 상위봉은 롤업하므로
    기본·권장은 1분 (env ``PERM_SUB_BASE_TF_MIN``).
    """
    ensure_permanent_subs_table(db)
    code = str(code or "").strip().upper()
    if not code:
        raise ValueError("code 필수")
    g_mt, g_ex, g_sym = classify_market(code)
    mt = (str(market_type).strip().upper() if market_type else "") or g_mt
    ex = (str(exchange).strip().upper() if exchange else "") or (g_ex if mt == g_mt else ("KRX" if mt == "KR" else "NASD"))
    sym = (str(symbol).strip().upper() if symbol else "") or g_sym or code
    try:
        from kis_trader.utils.env import get_env_int
        base_tf = int(get_env_int("PERM_SUB_BASE_TF_MIN", 1))
    except Exception:
        base_tf = 1
    if base_tf < 1:
        base_tf = 1
    # 요청값과 무관하게 기준봉으로 통일 (1분→롤업)
    try:
        _ = int(tf_min)
    except (TypeError, ValueError):
        pass
    tfv = base_tf
    raw = _core(db)
    raw.conn.execute(
        """
        INSERT INTO permanent_subscriptions
            (code, market_type, exchange, symbol, tf_min, enabled, note)
        VALUES (%s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
            market_type=VALUES(market_type), exchange=VALUES(exchange),
            symbol=VALUES(symbol), tf_min=VALUES(tf_min),
            enabled=VALUES(enabled), note=VALUES(note)
        """,
        [code, mt, ex, sym, tfv, 1 if enabled else 0, str(note or "")[:100]],
    )
    raw.conn.commit()


def normalize_all_tf_to_base(db: Any) -> int:
    """기존 영구구독 tf_min 을 기준봉(기본 1)으로 일괄 정리. 변경 건수 반환."""
    ensure_permanent_subs_table(db)
    try:
        from kis_trader.utils.env import get_env_int
        base_tf = int(get_env_int("PERM_SUB_BASE_TF_MIN", 1))
    except Exception:
        base_tf = 1
    if base_tf < 1:
        base_tf = 1
    raw = _core(db)
    cur = raw.conn.execute(
        "UPDATE permanent_subscriptions SET tf_min=%s WHERE tf_min<>%s",
        [base_tf, base_tf],
    )
    raw.conn.commit()
    try:
        return int(cur.rowcount or 0)
    except Exception:
        return 0

def remove_permanent_sub(db: Any, code: str) -> bool:
    ensure_permanent_subs_table(db)
    code = str(code or "").strip().upper()
    if not code:
        return False
    raw = _core(db)
    cur = raw.conn.execute("DELETE FROM permanent_subscriptions WHERE code=%s", [code])
    raw.conn.commit()
    try:
        return bool(cur.rowcount)
    except Exception:
        return True


def set_enabled(db: Any, code: str, enabled: bool) -> None:
    ensure_permanent_subs_table(db)
    raw = _core(db)
    raw.conn.execute(
        "UPDATE permanent_subscriptions SET enabled=%s WHERE code=%s",
        [1 if enabled else 0, str(code or "").strip().upper()],
    )
    raw.conn.commit()


def list_permanent_subs(db: Any, enabled_only: bool = False) -> List[Dict[str, Any]]:
    ensure_permanent_subs_table(db)
    raw = _core(db)
    sql = "SELECT code, market_type, exchange, symbol, tf_min, enabled, note FROM permanent_subscriptions"
    if enabled_only:
        sql += " WHERE enabled=1"
    sql += " ORDER BY market_type, code"
    cur = raw.conn.execute(sql)
    rows = cur.fetchall() if cur else []
    out: List[Dict[str, Any]] = []
    for r in rows:
        d = dict(r) if not isinstance(r, dict) else dict(r)
        d["enabled"] = int(d.get("enabled", 1))
        out.append(d)
    return out


def subscribe_master_enabled() -> bool:
    """마스터 스위치 — OFF면 행은 두고 구독만 안 함."""
    from kis_trader.utils.env import get_env_bool
    return get_env_bool("PERMANENT_SUBSCRIBE_ENABLED", True)


def enabled_code_set(db: Any, market_type: Optional[str] = None) -> set:
    """enabled=1 영구구독 코드. 마스터와 무관 (저장 가드는 이 집합)."""
    mt = str(market_type or "").strip().upper()
    out = set()
    for r in list_permanent_subs(db, enabled_only=True):
        if mt and str(r.get("market_type") or "KR").strip().upper() != mt:
            continue
        c = str(r.get("code") or "").strip()
        if c:
            out.add(c)
    return out


def subscribe_codes(db: Any, market_type: str) -> List[Dict[str, Any]]:
    """실제 WS 구독에 쓸 목록 = 마스터 ON ∧ enabled."""
    if not subscribe_master_enabled():
        return []
    return codes_by_market(db, market_type, enabled_only=True)


def should_persist_ls(code: str, perm_codes: Optional[set] = None) -> bool:
    """LS DB 적재: 영구구독 enabled 코드만 (후보 spill 제외)."""
    c = str(code or "").strip()
    if not c:
        return False
    if perm_codes is None:
        return False
    return c in perm_codes


def codes_by_market(db: Any, market_type: str, enabled_only: bool = True) -> List[Dict[str, Any]]:
    """특정 시장(KR/US) 영구구독 목록 (WS 배선용)."""
    mt = str(market_type or "").strip().upper()
    return [r for r in list_permanent_subs(db, enabled_only=enabled_only)
            if str(r.get("market_type", "KR")).strip().upper() == mt]


def migrate_env_codes(db: Any, env_csv: str) -> int:
    """
    기존 ``PERMANENT_WS_CODES`` 콤마 문자열 → 테이블 1회 이관.
    이미 있는 code 는 건드리지 않음(덮어쓰기 방지). 신규 삽입 건수 반환.
    """
    ensure_permanent_subs_table(db)
    existing = {r["code"] for r in list_permanent_subs(db)}
    n = 0
    for raw_code in str(env_csv or "").split(","):
        c = raw_code.strip().upper()
        if not c or c in existing:
            continue
        mt, ex, sym = classify_market(c)
        upsert_permanent_sub(db, c, mt, ex, sym, note="env 이관")
        existing.add(c)
        n += 1
    return n


def last_quotes_for_codes(
    db: Any,
    codes: List[str],
    tf_min: int = 1,
) -> Dict[str, Dict[str, Any]]:
    """영구구독 UI용 현재가 — REST/틱 전수스캔 없이 ``ws_candles`` 만.

    ※ 예전 ``ws_ticks`` ROW_NUMBER 전종목 스캔은 수십 초 지연 → UI 절대 금지.
       이 함수는 ``ws_candles`` 만 조회한다. ws_ticks 경로를 다시 넣지 말 것.

    반환 code → {
      price, price_src ('candle'|''), candle_time, tick_time(항상 ''),
      chg_pct (직전 봉 종가 대비 %), volume, updated_at
    }
    """
    from datetime import datetime as _dt

    out: Dict[str, Dict[str, Any]] = {}
    uniq = []
    seen = set()
    for c in codes or []:
        u = str(c or "").strip().upper()
        if not u or u in seen:
            continue
        seen.add(u)
        uniq.append(u)
        out[u] = {
            "price": 0.0,
            "price_src": "",
            "candle_time": "",
            "tick_time": "",
            "chg_pct": None,
            "volume": 0,
            "updated_at": "",
        }
    if not uniq:
        return out
    try:
        tf = int(tf_min) if int(tf_min) >= 1 else 1
    except (TypeError, ValueError):
        tf = 1
    # 미래 슬롯(이상봉) 제외 — 로컬 now + 2분 여유
    try:
        now_slot = (_dt.now()).strftime("%Y%m%d%H%M")
    except Exception:
        now_slot = "999999999999"
    raw = _core(db)
    for code in uniq:
        rows = []
        # KR 영구구독 시세는 ls_ws_candles 우선 (키움/KIS 슬롯 이관)
        if code.isdigit() and len(code) == 6:
            try:
                now_ls = ""
                if len(now_slot) >= 12:
                    now_ls = (
                        f"{now_slot[0:4]}-{now_slot[4:6]}-{now_slot[6:8]} "
                        f"{now_slot[8:10]}:{now_slot[10:12]}:00"
                    )
                ls_rows = raw.conn.execute(
                    """
                    SELECT datetime, open, close, volume, updated_at
                    FROM ls_ws_candles
                    WHERE code=%s AND tf_min=%s AND datetime<=%s
                    ORDER BY datetime DESC
                    LIMIT 2
                    """,
                    [code, tf, now_ls or "9999-12-31 23:59:00"],
                ).fetchall() or []
                for r in ls_rows:
                    d = dict(r) if not isinstance(r, dict) else dict(r)
                    dt = str(d.get("datetime") or "")
                    digits = "".join(ch for ch in dt if ch.isdigit())[:12]
                    d["candle_time"] = digits
                    rows.append(d)
            except Exception as e:
                logger.debug("last_quotes ls_ws_candles %s: %s", code, e)
                rows = []
        if not rows:
            try:
                rows = raw.conn.execute(
                    """
                    SELECT candle_time, open, close, volume, updated_at
                    FROM ws_candles
                    WHERE code=%s AND timeframe=%s AND candle_time<=%s
                    ORDER BY candle_time DESC
                    LIMIT 2
                    """,
                    [code, tf, now_slot],
                ).fetchall() or []
            except Exception as e:
                logger.debug("last_quotes candle %s: %s", code, e)
                rows = []
        if not rows:
            continue
        parsed: List[Dict[str, Any]] = []
        for r in rows:
            parsed.append(dict(r) if not isinstance(r, dict) else dict(r))
        last = parsed[0]
        try:
            px = float(last.get("close") or 0)
        except (TypeError, ValueError):
            px = 0.0
        if px <= 0:
            try:
                px = float(last.get("open") or 0)
            except (TypeError, ValueError):
                px = 0.0
        chg = None
        if len(parsed) >= 2 and px > 0:
            try:
                prev = float(parsed[1].get("close") or 0)
                if prev > 0:
                    chg = (px - prev) / prev * 100.0
            except (TypeError, ValueError, ZeroDivisionError):
                chg = None
        try:
            vol = int(last.get("volume") or 0)
        except (TypeError, ValueError):
            vol = 0
        out[code].update({
            "price": px,
            "price_src": "candle" if px > 0 else "",
            "candle_time": str(last.get("candle_time") or ""),
            "chg_pct": chg,
            "volume": vol,
            "updated_at": str(last.get("updated_at") or ""),
        })
    return out


def fill_ls_candles_from_kiwoom(
    db: Any,
    codes: List[str],
    *,
    n_bars: Optional[int] = None,
    tf_min: int = 1,
) -> Dict[str, Any]:
    """키움 ka10080 → ls_ws_candles INSERT IGNORE (구멍만). ws_candles 에 넣지 않음."""
    import random
    import time as _time

    from kis_trader.utils.env import get_env_bool, get_env_float, get_env_from_db, get_env_int
    from kis_trader.ws.kis_ws import get_kiwoom_candles_df
    from database import TradeDB

    tf = max(1, int(tf_min or 1))
    n_req = int(n_bars) if n_bars else int(get_env_int("PERM_LS_FILL_BARS", 500))
    sleep_lo = float(get_env_float("PERM_LS_FILL_SLEEP_MIN", 1.0))
    sleep_hi = float(get_env_float("PERM_LS_FILL_SLEEP_MAX", 3.0))
    if sleep_hi < sleep_lo:
        sleep_hi = sleep_lo

    force_real = get_env_bool("KIWOOM_WS_FORCE_REAL", True)
    is_mock = False if force_real else get_env_bool("KIS_MOCK", False)
    if is_mock:
        key = str(get_env_from_db("KIWOOM_APP_KEY_MOCK", "") or "").strip()
        sec = str(get_env_from_db("KIWOOM_APP_SECRET_MOCK", "") or "").strip()
    else:
        key = str(get_env_from_db("KIWOOM_APP_KEY_REAL", "") or "").strip()
        sec = str(get_env_from_db("KIWOOM_APP_SECRET_REAL", "") or "").strip()
    if not key or not sec:
        key = str(get_env_from_db("KIWOOM_APP_KEY", "") or "").strip()
        sec = str(get_env_from_db("KIWOOM_APP_SECRET", "") or "").strip()
    if not key or not sec:
        return {"ok": False, "error": "키움 API 키 없음", "codes": []}

    raw = _core(db)
    out_codes: List[Dict[str, Any]] = []
    for i, code in enumerate(codes or []):
        c = str(code or "").strip()
        if not (c.isdigit() and len(c) == 6):
            out_codes.append({"code": c, "ok": False, "error": "KR 6자리만", "inserted": 0})
            continue
        inserted = 0
        err = ""
        try:
            df = get_kiwoom_candles_df(c, tf, key, sec, is_mock=is_mock, n=n_req)
            if df is None or getattr(df, "empty", True):
                err = "빈응답"
            else:
                for _, rec in df.iterrows():
                    ct = str(rec.get("time") or "")[:12]
                    if len(ct) < 12:
                        continue
                    close = float(rec.get("close") or 0)
                    if close <= 0:
                        continue
                    dt = TradeDB._candle_time_to_ls_datetime(ct)
                    if not dt:
                        continue
                    candle = {
                        "datetime": dt,
                        "tf_min": tf,
                        "open": float(rec.get("open") or close),
                        "high": float(rec.get("high") or close),
                        "low": float(rec.get("low") or close),
                        "close": close,
                        "volume": float(rec.get("volume") or 0),
                        "tick_count": 0,
                    }
                    if raw.insert_ls_ws_candle_if_absent(code=c, candle=candle):
                        inserted += 1
        except Exception as e:
            err = str(e)
        out_codes.append({
            "code": c, "ok": not err, "error": err, "inserted": inserted,
        })
        if i + 1 < len(codes or []):
            _time.sleep(random.uniform(sleep_lo, sleep_hi))
    ok_n = sum(1 for x in out_codes if x.get("ok"))
    return {"ok": True, "codes": out_codes, "ok_n": ok_n}