kis_bot/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 60,
    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 = 60,
    enabled: bool = True,
    note: str = "",
) -> None:
    """코드 1건 등록/갱신 (code UNIQUE)."""
    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:
        tfv = int(tf_min)
    except (TypeError, ValueError):
        tfv = 60
    if tfv < 1:
        tfv = 60
    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 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 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