kis_bot/kis_trader/scan/dart_watchlist.py

#!/usr/bin/env python3
"""
kis_trader/scan/dart_watchlist.py — DART 임시 워치 (영구구독 아님)
"""
from __future__ import annotations

import datetime as dt
from typing import Any, Dict, List, Optional

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

logger = get_logger("kis_trader.dart_watch")


def _now() -> dt.datetime:
    return dt.datetime.now()


def _fmt(ts: dt.datetime) -> str:
    return ts.strftime("%Y-%m-%d %H:%M:%S")


def ensure_watchlist_table(db: Any) -> None:
    db.conn.execute(
        """
        CREATE TABLE IF NOT EXISTS dart_watchlist (
            stock_code VARCHAR(20) NOT NULL PRIMARY KEY,
            corp_name VARCHAR(100) NOT NULL DEFAULT '',
            rcept_no VARCHAR(32) NOT NULL DEFAULT '',
            report_nm VARCHAR(255) NOT NULL DEFAULT '',
            added_at VARCHAR(30) NOT NULL,
            expires_at VARCHAR(30) NOT NULL,
            enabled TINYINT NOT NULL DEFAULT 1,
            KEY idx_dart_watch_exp (expires_at)
        ) CHARACTER SET utf8mb4
        """
    )


def purge_expired(db: Any) -> int:
    ensure_watchlist_table(db)
    now = _fmt(_now())
    cur = db.conn.execute(
        "DELETE FROM dart_watchlist WHERE expires_at < %s OR enabled=0",
        (now,),
    )
    try:
        return int(cur.rowcount or 0)
    except Exception:
        return 0


def list_active_watch(db: Any) -> List[Dict[str, Any]]:
    ensure_watchlist_table(db)
    purge_expired(db)
    now = _fmt(_now())
    rows = db.conn.execute(
        """
        SELECT stock_code, corp_name, rcept_no, report_nm, added_at, expires_at
        FROM dart_watchlist
        WHERE enabled=1 AND expires_at >= %s
        ORDER BY added_at DESC
        """,
        (now,),
    ).fetchall()
    return [dict(r) for r in rows]


def watch_count(db: Any) -> int:
    return len(list_active_watch(db))


def upsert_watch(
    db: Any,
    *,
    stock_code: str,
    corp_name: str = "",
    rcept_no: str = "",
    report_nm: str = "",
    ttl_hours: Optional[int] = None,
    watch_max: Optional[int] = None,
) -> bool:
    """
    구독 스위치 ON일 때만 호출. 상한 초과 시 가장 오래된 것 제거 후 삽입.
    """
    if not get_env_bool("DART_SUBSCRIBE_ENABLED", False):
        return False
    code = (stock_code or "").strip()
    if not code:
        return False
    ensure_watchlist_table(db)
    purge_expired(db)
    ttl = int(ttl_hours if ttl_hours is not None else get_env_int("DART_WATCH_TTL_HOURS", 24))
    cap = int(watch_max if watch_max is not None else get_env_int("DART_WATCH_MAX", 15))
    now = _now()
    exp = now + dt.timedelta(hours=max(1, ttl))

    active = list_active_watch(db)
    if code not in {a["stock_code"] for a in active} and len(active) >= cap:
        # 가장 오래된 1건 제거
        oldest = sorted(active, key=lambda x: x.get("added_at") or "")[:1]
        for o in oldest:
            db.conn.execute(
                "DELETE FROM dart_watchlist WHERE stock_code=%s",
                (o["stock_code"],),
            )
            logger.info("DART 워치 상한 — 제거 %s", o["stock_code"])

    db.conn.execute(
        """
        INSERT INTO dart_watchlist
            (stock_code, corp_name, rcept_no, report_nm, added_at, expires_at, enabled)
        VALUES (%s, %s, %s, %s, %s, %s, 1)
        ON DUPLICATE KEY UPDATE
            corp_name=VALUES(corp_name),
            rcept_no=VALUES(rcept_no),
            report_nm=VALUES(report_nm),
            added_at=VALUES(added_at),
            expires_at=VALUES(expires_at),
            enabled=1
        """,
        (
            code[:20],
            (corp_name or "")[:100],
            (rcept_no or "")[:32],
            (report_nm or "")[:255],
            _fmt(now),
            _fmt(exp),
        ),
    )
    return True


def event_time_for_code(db: Any, code: str) -> Optional[str]:
    """워치/최신 공시 시각 → 봉키 YYYYMMDDHHMM 근사."""
    ensure_watchlist_table(db)
    code = (code or "").strip()
    row = db.conn.execute(
        """
        SELECT added_at, rcept_no FROM dart_watchlist
        WHERE stock_code=%s AND enabled=1
        ORDER BY added_at DESC LIMIT 1
        """,
        (code,),
    ).fetchone()
    if row and row.get("added_at"):
        s = str(row["added_at"]).replace("-", "").replace(":", "").replace(" ", "")
        if len(s) >= 12:
            return s[:12]
    # disclosures fallback
    d = db.conn.execute(
        """
        SELECT first_seen_at, rcept_dt FROM dart_disclosures
        WHERE stock_code=%s ORDER BY first_seen_at DESC LIMIT 1
        """,
        (code,),
    ).fetchone()
    if not d:
        return None
    if d.get("first_seen_at"):
        s = str(d["first_seen_at"]).replace("-", "").replace(":", "").replace(" ", "")
        if len(s) >= 12:
            return s[:12]
    rd = str(d.get("rcept_dt") or "")
    if len(rd) == 8:
        return rd + "0900"
    return None