# kiwwom_trader/scalping_bot.py
# 키움 REST API + 웹소켓 기반 스캘핑 봇. KIS kis_scalping_ver2와 동일한 로직·구조.
# 폴더를 /home/hoon 밑으로 옮겨도 config 경로 설정으로 동작.

import logging
import os
import time
from datetime import datetime as dt
from typing import Dict, List, Optional

# config는 run.py에서 먼저 import됨. 여기서는 ROOT만 참조
import config

logging.basicConfig(
    format="[%(asctime)s] %(message)s",
    datefmt="%H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger("KiwoomScalp")

# ── 의존성: kiwoom_rest_api (ROOT/.. 에 있거나 sys.path에 추가된 상태) ──
try:
    from kiwoom_rest_api.auth.token import TokenManager
    from kiwoom_rest_api.koreanstock.stockinfo import StockInfo
    from kiwoom_rest_api.koreanstock.chart import Chart
    from kiwoom_rest_api.koreanstock.order import Order
    from kiwoom_rest_api.koreanstock.rank_info import RankInfo
    from kiwoom_rest_api.koreanstock.account import Account
    from kiwoom_rest_api.koreanstock.market_condition import MarketCondition
    _KIWOOM_AVAILABLE = True
except ImportError as e:
    logger.warning("kiwoom_rest_api 미설치 또는 경로 미등록: %s", e)
    _KIWOOM_AVAILABLE = False

# ── kis_bot (database, scalping_engine 래핑 로직) ──
try:
    from database import TradeDB
    from scalping_logic import (
        load_defaults_from_db,
        check_buy_signal as logic_check_buy_signal,
    )
    _KIS_BOT_AVAILABLE = True
except ImportError as e:
    logger.warning(
        "kis_bot 경로 미등록(database, scalping_logic): %s → KIS_BOT_PATH 또는 동일 레포 내 실행 필요",
        e,
    )
    _KIS_BOT_AVAILABLE = False


def _get_env_from_db(db, key: str, default: str = ""):
    if not db:
        return os.environ.get(key, default)
    env_data = db.get_latest_env()
    if env_data and env_data.get("snapshot"):
        return env_data["snapshot"].get(key, default)
    return os.environ.get(key, default)


def _get_env_float(db, key: str, default: float):
    v = _get_env_from_db(db, key, str(default))
    if isinstance(v, str) and "#" in v:
        v = v.split("#")[0].strip()
    try:
        return float(v) if v else default
    except (ValueError, TypeError):
        return default


def _get_env_int(db, key: str, default: int):
    v = _get_env_from_db(db, key, str(default))
    if isinstance(v, str) and "#" in v:
        v = v.split("#")[0].strip()
    try:
        return int(v) if v else default
    except (ValueError, TypeError):
        return default


class ScalpingBotKiwoom:
    """
    키움 REST API로 주문·조회, 웹소켓으로 실시간 체결 수신 후 1분봉 집계.
    매수/매도 판단은 scalping_engine (kis_bot)과 동일.
    """

    def __init__(self):
        if not _KIWOOM_AVAILABLE:
            raise RuntimeError("kiwoom_rest_api를 사용할 수 없습니다. 경로 또는 설치를 확인하세요.")
        if not _KIS_BOT_AVAILABLE:
            raise RuntimeError("database, scalping_engine을 사용할 수 없습니다. KIS_BOT_PATH 또는 kis_bot 경로를 설정하세요.")

        self.db = TradeDB()
        self.token_manager = TokenManager()
        self.stock_info = StockInfo(token_manager=self.token_manager)
        self.chart = Chart(token_manager=self.token_manager)
        self.order = Order(token_manager=self.token_manager)
        self.rank = RankInfo(token_manager=self.token_manager)
        self.account = Account(token_manager=self.token_manager)
        self.market = MarketCondition(token_manager=self.token_manager)

        self.candle_tf = 1
        self.rsi_oversold = _get_env_float(self.db, "SCALP_RSI_OVERSOLD", 25.0)
        self.scalp_stop_loss_pct = _get_env_float(self.db, "SCALP_STOP_LOSS_PCT", 0.015)
        self.scalp_take_profit_pct = _get_env_float(self.db, "SCALP_TAKE_PROFIT_PCT", 0.015)
        self.scalp_min_drop_rate = _get_env_float(self.db, "SCALP_MIN_DROP_RATE", 0.015)
        self.max_stocks = _get_env_int(self.db, "SCALP_MAX_POSITIONS", 3)
        self._scan_engine_params = None
        self._candle_agg = None
        self._ws_client = None
        self._ws_thread = None

    def start_websocket_candles(self, codes: List[str]) -> None:
        """웹소켓 실시간 체결(0B) 구독 → 1분봉 집계. codes는 종목 코드 리스트 (최대 100개 권장)."""
        from ws_candle_agg import KiwoomCandleAggregator, build_ws_client_with_candle_callback
        token = self.token_manager.access_token
        if not token:
            logger.warning("웹소켓: 토큰 없음 → 스킵")
            return
        self._candle_agg = KiwoomCandleAggregator()
        client = build_ws_client_with_candle_callback(token, self._candle_agg)
        if client is None:
            return
        self._ws_client = client
        import asyncio
        async def run_ws():
            await client.start()
            await client.register_realtime(type_list=["0B"], item_list=codes if codes else [""])
            await client.run_forever()
        def thread_target():
            asyncio.run(run_ws())
        import threading
        self._ws_thread = threading.Thread(target=thread_target, daemon=True)
        self._ws_thread.start()
        logger.info("웹소켓 캔들 구독 시작: %d종목", len(codes))

    def _ensure_engine_params(self):
        if self._scan_engine_params is None:
            d = load_defaults_from_db()
            self._scan_engine_params = {
                **d,
                "rsi_oversold": self.rsi_oversold,
                "sl_pct": abs(self.scalp_stop_loss_pct),
                "tp_pct": self.scalp_take_profit_pct,
                "drop_rate": self.scalp_min_drop_rate,
            }

    def _norm_candle(self, c: dict) -> dict:
        """scalping_engine이 기대하는 키: candle_time, open, high, low, close, volume."""
        return {
            "candle_time": c.get("candle_time", ""),
            "open": float(c.get("open", 0)),
            "high": float(c.get("high", 0)),
            "low": float(c.get("low", 0)),
            "close": float(c.get("close", 0)),
            "volume": int(c.get("volume", 0)),
        }

    def get_candles(self, code: str, n: int = 50) -> List[dict]:
        """캔들 소스: 웹소켓 집계 → 없으면 DB ws_candles. scalping_engine 형식으로 반환."""
        if self._candle_agg:
            candles = self._candle_agg.get_candles(code, self.candle_tf, n=n)
            if candles:
                return [self._norm_candle(c) for c in candles]
        raw = self.db.get_ws_candles(code, self.candle_tf, limit=n, confirmed_only=True)
        return [self._norm_candle(c) for c in raw]

    def _get_ws_codes_from_universe_history(self, max_codes: int = 100) -> List[str]:
        """
        오늘자 유니버스 히스토리(target_candidates_history)에서
        슬롯 순서대로 종목 코드를 모아서 웹소켓 구독 리스트 생성.
        - slot_key(YYYYMMDDHHMM) 오름차순으로 flatten
        - 중복 제거, 최대 max_codes개까지
        - 히스토리가 비어 있으면 현재 target_candidates로 폴백
        """
        today_ymd = dt.now().strftime("%Y%m%d")
        codes: List[str] = []
        seen = set()
        try:
            universe_by_slot = self.db.get_universe_history_for_backtest(today_ymd, today_ymd)
            for slot_key in sorted(universe_by_slot.keys()):
                for code in universe_by_slot[slot_key]:
                    if not code or code in seen:
                        continue
                    seen.add(code)
                    codes.append(code)
                    if len(codes) >= max_codes:
                        break
                if len(codes) >= max_codes:
                    break
        except Exception as e:
            logger.warning("유니버스 히스토리 조회 실패 → 폴백 사용: %s", e)

        # 히스토리가 아직 없으면 현재 target_candidates 상위 N개로 폴백
        if not codes:
            try:
                current = self.db.get_target_candidates()
                for c in current:
                    code = c.get("code")
                    if not code or code in seen:
                        continue
                    seen.add(code)
                    codes.append(code)
                    if len(codes) >= max_codes:
                        break
            except Exception as e:
                logger.warning("폴백 후보 조회 실패: %s", e)

        return codes

    def check_buy_signal(self, code: str, name: str) -> Optional[dict]:
        """scalping_engine.check_buy_signal_live 사용 (KIS 봇과 동일)."""
        self._ensure_engine_params()
        candles = self.get_candles(code, n=50)
        if len(candles) < 5:
            return None
        today = dt.now().strftime("%Y%m%d")
        try:
            today_trades = self.db.get_trades_by_date(today)
            daily_cnt = len([t for t in today_trades if t.get("code") == code and str(t.get("strategy", "")).startswith("SCALP")])
        except Exception:
            daily_cnt = 0
        state = {"last_exit_dt": None, "daily_cnt": daily_cnt}
        reject_reason, reject_msg, sig = logic_check_buy_signal(candles, self._scan_engine_params, state)
        if reject_reason or sig is None:
            return None
        return sig

    def run(self):
        """메인 루프: 후보 조회 → 매수 신호 검사 → 주문 (구조만 연결, 실제 주문은 안전을 위해 로그만)."""
        logger.info("키움 스캘핑 봇 시작 (ROOT=%s)", config.ROOT)
        # 웹소켓 캔들: 후보가 있으면 0B 구독 (선택)
        if _get_env_from_db(self.db, "KIWOOM_WS_CANDLES", "0").lower() in ("1", "true", "yes"):
            try:
                codes = self._get_ws_codes_from_universe_history(max_codes=100)
                if codes:
                    self.start_websocket_candles(codes)
            except Exception as e:
                logger.warning("웹소켓 캔들 시작 스킵: %s", e)
        while True:
            try:
                candidates = self.db.get_target_candidates()
                if not candidates:
                    logger.info("후보 0건 → 대기")
                    time.sleep(60)
                    continue
                for c in candidates[: self.max_stocks * 2]:
                    code = c.get("code")
                    name = c.get("name", code)
                    if not code:
                        continue
                    sig = self.check_buy_signal(code, name)
                    if sig:
                        logger.info("[매수신호] %s %s RSI=%.1f (실주문은 주석 해제 후 사용)", name, code, sig.get("rsi"))
                        # self._place_buy(code, name, sig)  # 실제 주문 시 구현
                time.sleep(30)
            except KeyboardInterrupt:
                logger.info("종료")
                break
            except Exception as e:
                logger.exception("루프 오류: %s", e)
                time.sleep(60)
