#!/usr/bin/env python3
"""
LS 서버저장조건(동명) → ls_candidates_history 단독 수집기
========================================================

⚠️ 운영 기본 경로가 아님.
  ``LS_CONDITION_HISTORY_ENABLED=true`` 이면 **kis_trader_main** 이
  조건 매니저 + (FOLLOW 시) LS WS 를 기동한다.
  단독 수집기와 main 을 **동시에** 돌리면 AFR/유량 충돌 가능 → 둘 중 하나만.

실매(kis_trader_main) / 키움 조건검색 / target_candidates_history 를 **건드리지 않음**.

흐름:
  1) t1866 목록 → 이름(breakout/tail/momentum/scalp_re)으로 query_index 해결
  2) t1859 스냅샷 → 초기 유니버스
  3) t1860 실시간등록 → WS AFR(tr_key=sAlertNum) REG
  4) AFR 편입/이탈·주기 스냅샷 → ``ls_candidates_history`` INSERT

시세(US3) FOR 미러는 기본 OFF.
  --also-ticks 를 켜면 **현재 조건 유니버스 코드만** 천천히 US3 REG
  (조건명은 t1866 단계에서만 사용 — LS 시세 REG 는 코드/tr_key 만 받음).

사용:
  # 기존 LS 히스토리 비우고 수집 시작
  python3 scripts/collect_ls_condition_history.py --purge --run

  # 목록만
  python3 scripts/collect_ls_condition_history.py --list-only

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

import argparse
import importlib
import json
import logging
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))

# 기존 단독 테스트의 REST/WS 헬퍼 재사용 (실매 미의존)
_rt = importlib.import_module("_test_ls_condition_realtime")

# LS query_name → 전략 id (키움 CONDITION_*_KIWOOM_NAME 과 동일)
DEFAULT_NAME_TO_SID: Dict[str, str] = {
    "breakout": "BREAKOUT",
    "tail": "SHORT",
    "momentum": "MOMENTUM",
    "scalp_re": "SCALP",
}

DDL = """
CREATE TABLE IF NOT EXISTS ls_candidates_history (
    id BIGINT NOT NULL AUTO_INCREMENT PRIMARY KEY,
    slot_key VARCHAR(12) NOT NULL,
    scan_time VARCHAR(30) NOT NULL,
    code VARCHAR(20) NOT NULL,
    name VARCHAR(100) NOT NULL DEFAULT '',
    score DOUBLE NOT NULL DEFAULT 0,
    price DOUBLE NOT NULL DEFAULT 0,
    market CHAR(1) DEFAULT 'Q',
    sector VARCHAR(100) NULL,
    theme VARCHAR(100) NULL,
    strategy_id VARCHAR(32) NOT NULL DEFAULT '',
    event_time VARCHAR(30) NOT NULL DEFAULT '',
    query_name VARCHAR(64) NOT NULL DEFAULT '',
    query_index VARCHAR(32) NOT NULL DEFAULT '',
    source VARCHAR(16) NOT NULL DEFAULT 'ls_afr',
    INDEX idx_ls_cand_slot (slot_key),
    INDEX idx_ls_cand_sid_evt (strategy_id, event_time),
    INDEX idx_ls_cand_code (code),
    INDEX idx_ls_cand_scan (scan_time)
) CHARACTER SET utf8mb4
"""


def _setup_logger(log_path: Path) -> logging.Logger:
    logger = logging.getLogger("ls_cond_hist")
    logger.setLevel(logging.INFO)
    logger.handlers.clear()
    fmt = logging.Formatter("[%(asctime)s] %(message)s", datefmt="%H:%M:%S")
    sh = logging.StreamHandler(sys.stdout)
    sh.setFormatter(fmt)
    logger.addHandler(sh)
    log_path.parent.mkdir(parents=True, exist_ok=True)
    fh = logging.FileHandler(log_path, encoding="utf-8")
    fh.setFormatter(fmt)
    logger.addHandler(fh)
    return logger


def _ensure_table(db) -> None:
    db.conn.execute(DDL)


def _purge_ls_history(db, logger: logging.Logger) -> None:
    """기존 LS 쪽 히스토리만 비움 (키움 target_candidates_history / 실매 테이블 제외)."""
    for tbl in ("ls_universe_history", "ls_candidates_history"):
        try:
            db.conn.execute(f"CREATE TABLE IF NOT EXISTS {tbl} (id INT)")  # no-op if missing weird
        except Exception:
            pass
        try:
            n = db.conn.execute(f"SELECT COUNT(*) c FROM {tbl}").fetchone()["c"]
        except Exception as e:
            logger.warning("purge skip %s (없음?): %s", tbl, e)
            continue
        db.conn.execute(f"DELETE FROM {tbl}")
        logger.info("🗑️  purge %s rows_was=%s → 0", tbl, n)


def _slot_key_now() -> str:
    return datetime.now().strftime("%Y%m%d%H%M")


def _event_time_now() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def insert_ls_snapshot(
    db,
    *,
    strategy_id: str,
    query_name: str,
    query_index: str,
    items: List[Dict[str, Any]],
    event_time: Optional[str] = None,
    source: str = "ls_afr",
) -> int:
    """풀 유니버스 스냅샷 → ls_candidates_history (동일 event_time 재기록)."""
    et = event_time or _event_time_now()
    sk = et.replace("-", "").replace(":", "").replace(" ", "")[:12]
    sid = (strategy_id or "").strip().upper()
    qn = (query_name or "")[:64]
    qi = (query_index or "")[:32]
    inserted = 0
    with db.conn:
        db.conn.execute(
            "DELETE FROM ls_candidates_history "
            "WHERE strategy_id=%s AND event_time=%s",
            (sid, et),
        )
        for it in items:
            code = str(it.get("code") or "").strip()
            if not code:
                continue
            name = str(it.get("name") or code)[:100]
            try:
                price = float(it.get("price") or 0)
            except (TypeError, ValueError):
                price = 0.0
            db.conn.execute(
                """
                INSERT INTO ls_candidates_history
                    (slot_key, scan_time, code, name, score, price,
                     market, sector, theme, strategy_id, event_time,
                     query_name, query_index, source)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                """,
                (
                    sk, et, code, name, 0.0, price,
                    "Q", "", "", sid, et, qn, qi, source[:16],
                ),
            )
            inserted += 1
    return inserted


class ConditionState:
    def __init__(
        self,
        *,
        strategy_id: str,
        query_name: str,
        query_index: str,
        alert_num: str = "",
    ) -> None:
        self.strategy_id = strategy_id
        self.query_name = query_name
        self.query_index = query_index
        self.alert_num = alert_num
        self.codes: Dict[str, Dict[str, Any]] = {}  # code → {name, price}
        self.lock = threading.Lock()

    def set_from_snapshot(self, rows: List[Dict[str, Any]]) -> None:
        with self.lock:
            self.codes.clear()
            for r in rows:
                code = str(r.get("shcode") or r.get("code") or "").strip()
                if not code:
                    continue
                self.codes[code] = {
                    "code": code,
                    "name": str(r.get("hname") or r.get("name") or code)[:100],
                    "price": r.get("price") or 0,
                }

    def apply_afr(self, body: Dict[str, Any]) -> Tuple[str, Optional[str]]:
        """반환: (job, code) job=N/R/O."""
        job = str(body.get("gsJobFlag") or "").strip().upper()
        code = str(body.get("gsCode") or body.get("shcode") or "").strip()
        if not code:
            return job, None
        name = str(body.get("gsHname") or body.get("hname") or code)[:100]
        try:
            price = float(str(body.get("gsPrice") or body.get("price") or 0).replace(",", ""))
        except (TypeError, ValueError):
            price = 0.0
        with self.lock:
            if job in ("N", "R"):
                self.codes[code] = {"code": code, "name": name, "price": price}
            elif job == "O":
                self.codes.pop(code, None)
        return job, code

    def items(self) -> List[Dict[str, Any]]:
        with self.lock:
            return [dict(v) for v in self.codes.values()]


class AfrHistoryWatcher:
    """AFR 전용 WS — 프로토콜 ping OFF (클라 ping 대기 타임아웃 방지)."""

    def __init__(
        self,
        token: str,
        ws_url: str,
        logger: logging.Logger,
        *,
        by_alert: Dict[str, ConditionState],
        on_change,
        also_tick_codes: Optional[Set[str]] = None,
        tick_gap_ms: int = 200,
    ) -> None:
        import websocket

        self._websocket = websocket
        self.token = token
        self.ws_url = ws_url
        self.logger = logger
        self.by_alert = by_alert
        self.on_change = on_change
        self.also_tick_codes = set(also_tick_codes or ())
        self.tick_gap_ms = max(50, int(tick_gap_ms))
        self._ws: Any = None
        self._thread: Optional[threading.Thread] = None
        self.opened = threading.Event()
        self._stop = threading.Event()
        self.msg_count = 0
        self.afr_count = 0
        self._tick_sub: Set[str] = set()
        self._send_lock = threading.Lock()
        # 죽은 alert / rematch 전 미매칭 — 조용히 drop 하지 않음
        self._last_unmatched_afr_log = 0.0
        self._unmatched_afr_warn_sec = 30.0

    def start(self) -> None:
        self._ws = self._websocket.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._thread = threading.Thread(target=self._run, name="ls-afr-hist", daemon=True)
        self._thread.start()

    def _run(self) -> None:
        assert self._ws is not None
        # ping OFF — LS 가 클라 ping 에 pong 안 주는 케이스 회피
        self._ws.run_forever()

    def wait_open(self, timeout: float = 15.0) -> bool:
        return self.opened.wait(timeout=timeout)

    def _send(self, payload: dict) -> None:
        if self._ws is None:
            return
        with self._send_lock:
            self._ws.send(json.dumps(payload, ensure_ascii=False))

    def reg_afr(self, alert_num: str, *, tr_type: str = "3") -> None:
        self._send(
            {
                "header": {"token": self.token, "tr_type": tr_type},
                "body": {"tr_cd": "AFR", "tr_key": alert_num},
            }
        )
        self.logger.info(
            "AFR %s tr_key=%r (조건명→alert 매핑 후 키만 전송)",
            "REG" if tr_type == "3" else "UNREG",
            alert_num,
        )

    def reg_us3(self, code: str, *, tr_type: str = "3") -> None:
        """시세 미러(옵션). tr_key = U+6자리 — 이름은 WS REG 에 없음."""
        key = code if code.startswith("U") else f"U{code}"
        self._send(
            {
                "header": {"token": self.token, "tr_type": tr_type},
                "body": {"tr_cd": "US3", "tr_key": key},
            }
        )

    def sync_ticks(self, codes: Set[str]) -> None:
        if not self.also_tick_codes and not codes:
            return
        # also_tick_codes 가 비어있으면 유니버스 전체, 아니면 교집합만
        want = set(codes)
        if self.also_tick_codes:
            # True sentinel: empty set means "all current universe" when flag set
            pass
        add = want - self._tick_sub
        rem = self._tick_sub - want
        for c in sorted(rem):
            self.reg_us3(c, tr_type="4")
            self._tick_sub.discard(c)
            time.sleep(self.tick_gap_ms / 1000.0)
        for c in sorted(add):
            self.reg_us3(c, tr_type="3")
            self._tick_sub.add(c)
            time.sleep(self.tick_gap_ms / 1000.0)

    def stop(self) -> None:
        self._stop.set()
        if self._ws is not None:
            try:
                self._ws.close()
            except Exception:
                pass

    def _on_open(self, _ws: Any) -> None:
        self.logger.info("WS OPEN %s ping=off", self.ws_url)
        self.opened.set()
        # 재연결 시 AFR 재 REG (갈아끼우기 = alert 몇 개만)
        for alert, st in self.by_alert.items():
            if alert:
                try:
                    self.reg_afr(alert, tr_type="3")
                    time.sleep(0.3)
                except Exception as e:
                    self.logger.warning("AFR re-REG 실패 %s: %s", alert, e)

    def _on_close(self, _ws: Any, status: Any, msg: Any) -> None:
        self.opened.clear()
        self.logger.warning("WS CLOSE status=%s msg=%s", status, msg)

    def _on_error(self, _ws: Any, err: Any) -> None:
        self.logger.warning("WS ERROR %s", err)

    def _on_message(self, _ws: Any, message: Any) -> None:
        try:
            data = json.loads(message) if isinstance(message, str) else message
        except Exception:
            return
        self.msg_count += 1
        header = data.get("header") or {}
        body = data.get("body") or {}
        if not isinstance(body, dict):
            return
        tr_cd = str(header.get("tr_cd") or "")
        tr_key = str(header.get("tr_key") or body.get("tr_key") or "")
        if tr_cd != "AFR" and not body.get("gsCode"):
            return
        st = self.by_alert.get(tr_key)
        if st is None and len(self.by_alert) == 1:
            st = next(iter(self.by_alert.values()))
        if st is None:
            # REG ACK 등 tr_key 빈 프레임은 무시 (미매칭 경고 폭주 방지)
            if not str(tr_key or "").strip():
                return
            now = time.time()
            if now - float(self._last_unmatched_afr_log or 0.0) >= float(
                self._unmatched_afr_warn_sec or 30.0
            ):
                self._last_unmatched_afr_log = now
                self.logger.warning(
                    "⚠️ AFR 미매칭 tr_key=%r (등록 alert=%s) — "
                    "서버 조건 CRUD 후 죽은 키일 수 있음. "
                    "LS rematch(t1866) 대기/확인",
                    tr_key,
                    list(self.by_alert.keys()),
                )
            return
        job, code = st.apply_afr(body)
        if not code:
            return
        self.afr_count += 1
        self.logger.info(
            "AFR %s %s sid=%s name=%s px=%s alert=%s",
            _rt.AFR_JOB_KO.get(job, job or "?"),
            code,
            st.strategy_id,
            st.query_name,
            (body.get("gsPrice") or body.get("price")),
            tr_key,
        )
        try:
            # job/code/body 는 진입·재진입·이탈 이벤트 DB 적재용 (구 콜백은 st 만)
            try:
                self.on_change(st, job=job, code=code, body=body)
            except TypeError:
                self.on_change(st)
        except Exception as e:
            self.logger.warning("on_change: %s", e)


def _parse_names(s: str) -> List[str]:
    return [x.strip() for x in (s or "").split(",") if x.strip()]


def main() -> int:
    ap = argparse.ArgumentParser(description="LS 조건식 → ls_candidates_history 단독 수집")
    ap.add_argument("--list-only", action="store_true")
    ap.add_argument("--purge", action="store_true", help="ls_universe/ls_candidates history DELETE")
    ap.add_argument("--run", action="store_true", help="수집 루프 실행")
    ap.add_argument(
        "--names",
        default="breakout,tail,momentum,scalp_re",
        help="LS query_name 콤마목록 (기본=키움 동명)",
    )
    ap.add_argument(
        "--flush-sec",
        type=int,
        default=60,
        help="변동 없어도 N초마다 풀스냅샷 재저장 (기본 60)",
    )
    ap.add_argument(
        "--also-ticks",
        action="store_true",
        help="조건 유니버스 코드만 US3 시세 REG (느린 gap, 기본 OFF)",
    )
    ap.add_argument("--tick-gap-ms", type=int, default=200)
    ap.add_argument("--mock", action="store_true")
    ap.add_argument("--user-id", default="")
    args = ap.parse_args()

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_path = ROOT / "logs" / f"ls_condition_history_{ts}.log"
    logger = _setup_logger(log_path)
    logger.info("log → %s", log_path)
    logger.info(
        "purge=%s run=%s names=%r also_ticks=%s flush=%ds",
        args.purge, args.run, args.names, args.also_ticks, args.flush_sec,
    )

    from database import TradeDB

    db = TradeDB()
    _ensure_table(db)
    if args.purge:
        _purge_ls_history(db, logger)

    user_id = (args.user_id or "").strip()
    if not user_id:
        try:
            from kis_trader.utils.env import get_env_from_db

            user_id = (get_env_from_db("LS_USER_ID", "") or "").strip()
        except Exception:
            user_id = ""
    if not user_id:
        user_id = (os.environ.get("LS_USER_ID") or "").strip()
    if not user_id:
        logger.error("LS_USER_ID 필요")
        return 1

    app_key, app_secret = _rt.load_ls_creds(use_mock=bool(args.mock))
    if not app_key or not app_secret:
        logger.error("LS 앱키 없음")
        return 1
    token = _rt.fetch_access_token(app_key, app_secret)
    logger.info("token OK user_id=%s", user_id)

    rows = _rt.t1866_list_conditions(token, user_id, logger=logger)
    logger.info("서버저장조건 %d건", len(rows))
    for i, r in enumerate(rows):
        logger.info(
            "  [%d] index=%s name=%s group=%s",
            i, r["query_index"], r["query_name"], r["group_name"],
        )
    if args.list_only and not args.run:
        return 0

    want_names = _parse_names(args.names)
    states: List[ConditionState] = []
    for nm in want_names:
        hit = _rt._resolve_query(rows, name=nm, query_index="")
        if not hit or not hit.get("query_index"):
            logger.error("조건명 미매칭: %r", nm)
            return 3
        sid = DEFAULT_NAME_TO_SID.get(nm.lower(), nm.upper())
        st = ConditionState(
            strategy_id=sid,
            query_name=str(hit.get("query_name") or nm),
            query_index=str(hit["query_index"]),
        )
        states.append(st)
        logger.info(
            "매핑 name=%s → sid=%s query_index=%s",
            st.query_name, st.strategy_id, st.query_index,
        )

    # 스냅샷
    for st in states:
        snap = _rt.t1859_snapshot(token, st.query_index, logger=logger)
        st.set_from_snapshot(snap)
        n = insert_ls_snapshot(
            db,
            strategy_id=st.strategy_id,
            query_name=st.query_name,
            query_index=st.query_index,
            items=st.items(),
            source="ls_t1859",
        )
        logger.info(
            "📸 [%s/%s] snap=%d → ls_candidates_history +%d",
            st.strategy_id, st.query_name, len(st.items()), n,
        )

    if not args.run:
        logger.info("스냅샷만 완료 (--run 없음). log=%s", log_path)
        return 0

    # 실시간 등록
    by_alert: Dict[str, ConditionState] = {}
    for st in states:
        ob = _rt.t1860_realtime(
            token, st.query_index, flag="E", alert_num="", logger=logger,
        )
        logger.info("t1860[%s] %s", st.query_name, json.dumps(ob, ensure_ascii=False))
        if str(ob.get("sResultFlag") or "").strip() != "S":
            logger.error("t1860 실패 name=%s", st.query_name)
            return 5
        alert = str(ob.get("sAlertNum") or "").strip()
        if not alert:
            logger.error("sAlertNum 없음 name=%s", st.query_name)
            return 5
        st.alert_num = alert
        by_alert[alert] = st
        logger.info(
            "✅ 등록 name=%s sid=%s alert=%s (WS에는 alert만 REG — 이름은 로그/DB용)",
            st.query_name, st.strategy_id, alert,
        )

    save_lock = threading.Lock()
    last_flush = time.monotonic()

    def on_change(st: ConditionState) -> None:
        nonlocal last_flush
        with save_lock:
            n = insert_ls_snapshot(
                db,
                strategy_id=st.strategy_id,
                query_name=st.query_name,
                query_index=st.query_index,
                items=st.items(),
                source="ls_afr",
            )
            last_flush = time.monotonic()
            logger.info(
                "💾 [%s] universe=%d saved=%d",
                st.strategy_id, len(st.items()), n,
            )

    ws_url = _rt.LS_WS_MOCK if args.mock else _rt.LS_WS_REAL
    watcher = AfrHistoryWatcher(
        token,
        ws_url,
        logger,
        by_alert=by_alert,
        on_change=on_change,
        also_tick_codes=set() if args.also_ticks else None,
        tick_gap_ms=args.tick_gap_ms,
    )
    # also_ticks: 빈 set 이면 "유니버스 전부" 의미로 sync — flag만 True 일 때
    if args.also_ticks:
        watcher.also_tick_codes = set()  # unused marker; sync uses universe
    # AFR REG 는 on_open 에서만 (재연결 복구 포함) — 여기서 또내면 이중 REG
    watcher.start()
    if not watcher.wait_open(20.0):
        logger.error("WS OPEN 타임아웃")
        watcher.stop()
        return 4
    time.sleep(0.5 + 0.35 * max(1, len(by_alert)))

    if args.also_ticks:
        all_codes: Set[str] = set()
        for st in states:
            all_codes |= {it["code"] for it in st.items()}
        logger.info("US3 시세 미러 시작 codes=%d gap=%dms", len(all_codes), args.tick_gap_ms)
        watcher.sync_ticks(all_codes)

    logger.info("수집 루프 시작 (Ctrl+C 종료). flush=%ds", args.flush_sec)
    try:
        while not watcher._stop.is_set():
            time.sleep(1.0)
            if time.monotonic() - last_flush >= max(10, int(args.flush_sec)):
                with save_lock:
                    for st in states:
                        n = insert_ls_snapshot(
                            db,
                            strategy_id=st.strategy_id,
                            query_name=st.query_name,
                            query_index=st.query_index,
                            items=st.items(),
                            source="ls_flush",
                        )
                        logger.info(
                            "⏱️ flush [%s/%s] n=%d univ=%d",
                            st.strategy_id, st.query_name, n, len(st.items()),
                        )
                    last_flush = time.monotonic()
                if args.also_ticks:
                    all_codes = set()
                    for st in states:
                        all_codes |= {it["code"] for it in st.items()}
                    watcher.sync_ticks(all_codes)
    except KeyboardInterrupt:
        logger.info("KeyboardInterrupt → 정리")
    finally:
        for st in states:
            if not st.alert_num:
                continue
            try:
                watcher.reg_afr(st.alert_num, tr_type="4")
            except Exception as e:
                logger.warning("AFR UNREG: %s", e)
            try:
                _rt.t1860_realtime(
                    token,
                    st.query_index,
                    flag="D",
                    alert_num=st.alert_num,
                    logger=logger,
                )
            except Exception as e:
                logger.warning("t1860 D: %s", e)
        watcher.stop()
        logger.info(
            "종료 afr_events≈%d ws_msg=%d log=%s",
            watcher.afr_count, watcher.msg_count, log_path,
        )
    return 0


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