#!/usr/bin/env python3
"""
kis_token_manager.py — KIS 실전/모의 토큰 통합 자동 관리
==========================================================
문제: 모의 모드로만 봇을 돌리면 실전 토큰이 갱신되지 않음.
     → 실전 토큰 만료 → WebSocket/홀딩봇 REST 호출 실패 → 거래 없음
     → 20h 표류 타이머 + 장중 10분 선제 갱신으로 발급 시각이 장중에 밀림

해결:
  - 매일 장전(systemd OnCalendar 08:00)에 ensure_both_tokens()
  - 실전+모의 둘 다: **오늘 세션을 못 덮거나 이미 만료일 때만** 발급
  - 장중 get_token(): **완전 만료(±60초)일 때만** 비상 발급 (선제 SMS 금지)

KIS 접근토큰(/oauth2/tokenP) 한도 (우선 준수):
  - 유효 약 24시간
  - 1일 1회 발급 원칙 (유효기간 내 잦은 발급 → 이용 제한 가능)
  - EGW00133: 1분당 1회
  ※ WS approval_key 의 6h 재발급 가드는 별도(kis_approval_manager) — 여기서 건드리지 않음

사용 예시 (각 봇 시작부):
    from kis_token_manager import ensure_both_tokens
    ensure_both_tokens()   # 실전+모의: 세션 미커버분만 갱신

CLI:
    python3 kis_token_manager.py           # 상태 확인 + 필요 시만 갱신
    python3 kis_token_manager.py --refresh # 강제 갱신 시도 (한도 주의)
"""

import json
import logging
import os
import threading
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional

import requests

logger = logging.getLogger(__name__)

ROOT            = Path(__file__).parent
CACHE_MOCK      = ROOT / ".kis_token_cache_mock.json"
CACHE_REAL      = ROOT / ".kis_token_cache_real.json"
LOCK_FILE       = ROOT / ".kis_token_manager.lock"
# 장중 비상: 만료 N초 전부터만 재발급 시도 (옛 10분 선제 = 장중 카톡 유발 → 폐기)
EXPIRE_SAFETY_SEC = 60
# 오늘 정규장+버퍼를 덮는지 판정 (HHMM). 이 시각 전에 만료면 장전 발급.
SESSION_END_HHMM  = 1600
LOCK_TIMEOUT_S  = 60     # 잠금 최대 대기 시간(초)


# ─────────────────────────────────────────────────────────────────────────────
# 내부 유틸
# ─────────────────────────────────────────────────────────────────────────────
def _load_env() -> dict:
    """DB에서 KIS 키 로드 — 인증키는 env_auth_config 우선(get_auth_env)."""
    try:
        from database import TradeDB
        db = TradeDB()
        snap = db.get_merged_env_snapshot() or {}
        snap.update(db.get_auth_env())
        db.close()
        return snap
    except Exception as e:
        logger.warning(f"DB env 로드 실패: {e}")
        return {}


def _parse_expired(s: str):
    """만료 시간 문자열 → datetime. 실패 시 None."""
    for fmt in ("%Y-%m-%d %H:%M:%S", "%Y%m%d%H%M%S"):
        try:
            return datetime.strptime(str(s).strip(), fmt)
        except Exception:
            pass
    return None


def _session_end_hhmm() -> int:
    """오늘 세션 커버 목표 HHMM (기본 1600). OS env 로만 오버라이드."""
    raw = (os.environ.get("KIS_TOKEN_SESSION_END_HHMM", "") or "").strip()
    if not raw:
        return SESSION_END_HHMM
    try:
        v = int(raw)
        if 0 <= v <= 2359:
            return v
    except (TypeError, ValueError):
        pass
    return SESSION_END_HHMM


def _session_cover_deadline(now: Optional[datetime] = None) -> datetime:
    """당일 세션 커버 데드라인 (시:분)."""
    now = now or datetime.now()
    hhmm = _session_end_hhmm()
    return now.replace(
        hour=hhmm // 100,
        minute=hhmm % 100,
        second=0,
        microsecond=0,
    )


def token_covers_session(
    exp_dt: Optional[datetime],
    now: Optional[datetime] = None,
) -> bool:
    """
    접근토큰이 '오늘 운용'을 덮는지.
    - 이미 만료 → False (발급 필요)
    - 지금이 세션 종료 이전이면: 만료시각 >= 세션종료 이어야 True
    - 세션 종료 이후(장후)면: 아직 만료 전이면 True (불필요 재발급 금지)
    """
    now = now or datetime.now()
    if exp_dt is None:
        return False
    if now >= exp_dt:
        return False
    deadline = _session_cover_deadline(now)
    if now >= deadline:
        return True
    return exp_dt >= deadline


def get_token_status(is_mock: bool, current_app_key: str = None) -> dict:
    """
    캐시 파일 상태 반환.
    반환: valid=세션커버(ensure 재사용 기준), usable=만료 전 API 사용 가능
    """
    cache_path = CACHE_MOCK if is_mock else CACHE_REAL
    if not cache_path.exists():
        return {
            "valid": False,
            "usable": False,
            "token": "",
            "expires": "파일없음",
            "expires_in_h": -999,
        }
    try:
        cache      = json.loads(cache_path.read_text(encoding="utf-8"))
        token      = cache.get("access_token", "")
        expired_s  = cache.get("access_token_token_expired", "")
        exp_dt     = _parse_expired(expired_s)
        app_key_prefix = cache.get("app_key_prefix", "")

        if current_app_key and app_key_prefix:
            if not current_app_key.startswith(app_key_prefix):
                return {
                    "valid": False,
                    "usable": False,
                    "token": "",
                    "expires": "앱키변경됨",
                    "expires_in_h": -999,
                }

        if not token or exp_dt is None:
            return {
                "valid": False,
                "usable": False,
                "token": "",
                "expires": expired_s,
                "expires_in_h": -999,
            }
        now        = datetime.now()
        expires_in = (exp_dt - now).total_seconds() / 3600
        usable     = now < exp_dt
        covers     = token_covers_session(exp_dt, now)
        return {
            "valid":        covers,
            "usable":       usable,
            "token":        token[:12] + "…",
            "expires":      expired_s,
            "expires_in_h": round(expires_in, 1),
        }
    except Exception as e:
        return {
            "valid": False,
            "usable": False,
            "token": "",
            "expires": f"읽기오류: {e}",
            "expires_in_h": -999,
        }


def _acquire_lock() -> bool:
    """파일 잠금 획득 (중복 발급 방지). 성공 시 True."""
    deadline = time.time() + LOCK_TIMEOUT_S
    while time.time() < deadline:
        try:
            # O_CREAT | O_EXCL: 파일이 없을 때만 생성 (원자적)
            fd = os.open(str(LOCK_FILE), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            os.write(fd, str(os.getpid()).encode())
            os.close(fd)
            return True
        except FileExistsError:
            # 기존 잠금이 5분 이상 됐으면 강제 해제 (좀비 잠금)
            try:
                mtime = LOCK_FILE.stat().st_mtime
                if time.time() - mtime > 300:
                    LOCK_FILE.unlink(missing_ok=True)
                    continue
            except Exception:
                pass
            time.sleep(2)
        except Exception as e:
            logger.warning(f"잠금 획득 실패: {e}")
            return False
    logger.warning("잠금 획득 타임아웃 → 갱신 건너뜀")
    return False


def _release_lock():
    """파일 잠금 해제"""
    try:
        LOCK_FILE.unlink(missing_ok=True)
    except Exception:
        pass


# ─────────────────────────────────────────────────────────────────────────────
# 발급 시도 백오프 (분당 1회 제한 / 네트워크 오류 누적 폭주 방지)
# ─────────────────────────────────────────────────────────────────────────────
# KIS 가 EGW00133 (1분당 1회) 로 거부했을 때, 다음 시도까지 강제로 기다리는
# 시간(초). 이 백오프가 없으면 토큰 만료 시점에 호출자(전략·헬스체크·봇 등)
# 가 매번 _issue_token 을 호출 → KIS 정책 위반 패턴으로 분류될 위험.
RATE_LIMIT_BACKOFF_S    = 65    # EGW00133 받으면 65초 동안 재시도 금지
GENERIC_FAIL_BACKOFF_S  = 10    # 그 외 발급 실패 (네트워크/타임아웃 등)

# is_mock(bool) → 다음 시도 가능 시각 (epoch seconds)
_next_issue_at: dict = {False: 0.0, True: 0.0}
_next_issue_lock     = threading.Lock()


def _can_issue_now(is_mock: bool) -> tuple:
    """
    백오프 시간이 지났는지 확인.

    Returns:
        (allowed, wait_s) — allowed=False 면 wait_s 초 더 기다려야 함.
    """
    with _next_issue_lock:
        now      = time.time()
        next_at  = _next_issue_at.get(is_mock, 0.0)
        if now < next_at:
            return False, next_at - now
        return True, 0.0


def _set_backoff(is_mock: bool, seconds: float) -> None:
    """다음 발급 시도까지 ``seconds`` 초 대기."""
    with _next_issue_lock:
        _next_issue_at[is_mock] = time.time() + seconds


def _issue_token(app_key: str, app_secret: str, is_mock: bool) -> bool:
    """
    KIS /oauth2/tokenP 엔드포인트로 새 토큰 발급 후 캐시 저장.
    성공 시 True, 실패 시 False.

    ※ KIS 정책: 1일 1회 원칙 + 분당 1회 제한 (EGW00133).
       ensure_token 이 세션 미커버/만료일 때만 여기로 온다.
       발급 실패 시 백오프(`_set_backoff`)로 호출자 폭주 방지.
    """
    # 백오프 미경과 → 즉시 실패 (호출자에게 재시도 부담 떠넘기기 금지)
    allowed, wait_s = _can_issue_now(is_mock)
    if not allowed:
        logger.warning(
            "⏸ %s 토큰 발급 백오프 중 (%.1fs 남음) → 발급 시도 건너뜀",
            "모의" if is_mock else "실전", wait_s,
        )
        return False

    base_url   = (
        "https://openapivts.koreainvestment.com:29443"
        if is_mock else
        "https://openapi.koreainvestment.com:9443"
    )
    cache_path = CACHE_MOCK if is_mock else CACHE_REAL
    mode       = "모의" if is_mock else "실전"

    try:
        resp = requests.post(
            f"{base_url}/oauth2/tokenP",
            json={"grant_type": "client_credentials",
                  "appkey": app_key, "appsecret": app_secret},
            timeout=15,
        )
        data  = resp.json()
        token = data.get("access_token", "")
        exp   = data.get("access_token_token_expired", "")

        if not token:
            err = data.get("message") or data.get("error_description") or str(data)
            logger.error(f"❌ {mode} 토큰 발급 실패: {err}")
            # EGW00133: 1분당 1회 제한 → 65초 강제 백오프 (호출자 폭주 차단)
            if "EGW00133" in str(data):
                _set_backoff(is_mock, RATE_LIMIT_BACKOFF_S)
                logger.warning(
                    "⚠️ KIS 분당 발급 제한 (EGW00133). %ds 동안 재시도 차단.",
                    RATE_LIMIT_BACKOFF_S,
                )
            else:
                _set_backoff(is_mock, GENERIC_FAIL_BACKOFF_S)
            return False

        cache_path.write_text(
            json.dumps({
                "access_token":                 token,
                "access_token_token_expired":   exp,
                "mock":                         is_mock,
                "issued_at":                    datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "app_key_prefix":               app_key[:8] if app_key else "",
            }, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        # 성공 시 백오프 해제 (다음 만료 때 즉시 시도 가능)
        _set_backoff(is_mock, 0)
        logger.info(f"✅ {mode} 토큰 발급 완료 | 만료: {exp} | 앞12자: {token[:12]}…")
        return True

    except Exception as e:
        logger.error(f"❌ {mode} 토큰 발급 예외: {e}")
        _set_backoff(is_mock, GENERIC_FAIL_BACKOFF_S)
        return False


# ─────────────────────────────────────────────────────────────────────────────
# 공개 API
# ─────────────────────────────────────────────────────────────────────────────
def ensure_token(is_mock: bool, env: dict = None) -> bool:
    """
    단일 모드(실전/모의) 토큰: 오늘 세션을 덮으면 재사용, 아니면만 발급.
    1일 1회 원칙 — 충분하면 tokenP 호출 없음.
    """
    if env is None:
        env = _load_env()

    key_suffix = "MOCK" if is_mock else "REAL"
    app_key    = str(env.get(f"KIS_APP_KEY_{key_suffix}", "") or "").strip()
    app_secret = str(env.get(f"KIS_APP_SECRET_{key_suffix}", "") or "").strip()

    status = get_token_status(is_mock, current_app_key=app_key)
    mode   = "모의" if is_mock else "실전"

    if status["valid"]:
        logger.info(
            f"🔑 {mode} 토큰 세션커버 OK ({status['expires_in_h']:.1f}h 남음, "
            f"만료 {status['expires']}) → 재사용 (발급 안 함)"
        )
        return True

    logger.info(
        f"⏰ {mode} 토큰 세션미커버/만료 → 갱신 시도 "
        f"(남은시간: {status['expires_in_h']:.1f}h, 만료 {status['expires']})"
    )

    if env is None:
        env = _load_env()

    key_suffix = "MOCK" if is_mock else "REAL"
    app_key    = str(env.get(f"KIS_APP_KEY_{key_suffix}", "") or "").strip()
    app_secret = str(env.get(f"KIS_APP_SECRET_{key_suffix}", "") or "").strip()

    if not app_key or not app_secret:
        logger.warning(f"⚠️ {mode} 키 없음 (KIS_APP_KEY_{key_suffix}) → 발급 불가")
        return False

    if not _acquire_lock():
        return False
    try:
        # 잠금 획득 후 다시 확인 (다른 프로세스가 갱신했을 수 있음)
        status = get_token_status(is_mock, current_app_key=app_key)
        if status["valid"]:
            logger.info(f"🔑 {mode} 토큰 이미 갱신됨 (다른 프로세스) → 재사용")
            return True
        return _issue_token(app_key, app_secret, is_mock)
    finally:
        _release_lock()


def ensure_both_tokens(env: dict = None) -> dict:
    """
    실전 + 모의 토큰 모두 세션커버 유지.

    어느 모드로 봇을 실행하든 두 토큰을 동시에 점검.
    세션 미커버/만료분만 발급 → KIS 1일 1회 원칙 준수.

    반환: {"real": True/False, "mock": True/False}
    """
    if env is None:
        env = _load_env()

    results = {}
    # 실전 먼저 (더 중요, 잠금 경쟁 최소화)
    for is_mock in [False, True]:
        mode             = "mock" if is_mock else "real"
        results[mode]    = ensure_token(is_mock, env)

    ok_str = " | ".join(
        f"{'모의' if k == 'mock' else '실전'}={'✅' if v else '❌'}"
        for k, v in results.items()
    )
    logger.info(f"🔄 토큰 상태: {ok_str}")
    return results


# ─────────────────────────────────────────────────────────────────────────────
# KisTokenManager — KiwoomTokenManager 와 동일한 get_token() 인터페이스
#
# 동작 원리:
#   1. 메모리 캐시 (_token, _expiry) — 파일 읽기 최소화
#   2. JSON 파일 캐시 — 크로스 프로세스 공유 + 재시작 후 토큰 유지
#   3. 만료(±60초) 시에만 비상 발급 — 장중 10분 선제 발급 금지
#
# KIS 에는 토큰 상태 조회 API 가 없음 → 발급 시 받은 expires_dt 를
# 서버 현재시간과 비교해서 판단 (JSON 파일에 저장된 날짜 사용)
#
# 봇에서 사용:
#   from kis_token_manager import KisTokenManager
#   token = KisTokenManager.instance(is_mock).get_token()  # 매 API 호출 전
# ─────────────────────────────────────────────────────────────────────────────
_km_instances: dict = {}          # is_mock(bool) → KisTokenManager 싱글톤
_km_instances_lock  = threading.Lock()


class KisTokenManager:
    """
    KIS 토큰 싱글톤 매니저.
    - get_token(): 유효하면 바로 반환, 완전 만료 시에만 비상 갱신
    - kiwoom_rest_api TokenManager / KiwoomTokenManager 와 동일한 인터페이스
    """

    @classmethod
    def instance(cls, is_mock: bool) -> "KisTokenManager":
        """프로세스당 실전/모의 각 1개 싱글톤 반환"""
        with _km_instances_lock:
            if is_mock not in _km_instances:
                _km_instances[is_mock] = cls(is_mock)
            return _km_instances[is_mock]

    def __init__(self, is_mock: bool):
        self._is_mock    = is_mock
        self._mode_str   = "모의" if is_mock else "실전"
        self._cache_path = CACHE_MOCK if is_mock else CACHE_REAL
        self._lock       = threading.Lock()
        self._token:  Optional[str]      = None
        self._expiry: Optional[datetime] = None
        self._app_key_prefix: Optional[str] = None
        
        env = _load_env()
        if env:
            suffix = "MOCK" if is_mock else "REAL"
            key = str(env.get(f"KIS_APP_KEY_{suffix}", "")).strip()
            if key:
                self._app_key_prefix = key[:8]
                
        self._load_from_file()  # 재시작 후에도 기존 토큰 재사용

    # ── 내부 ──────────────────────────────────────────────────────
    def _load_from_file(self) -> None:
        """JSON 파일 → 메모리 로드"""
        if not self._cache_path.exists():
            return
        try:
            data   = json.loads(self._cache_path.read_text(encoding="utf-8"))
            token  = data.get("access_token", "")
            exp_dt = _parse_expired(data.get("access_token_token_expired", ""))
            
            # 앱키 변경 감지: 캐시된 app_key_prefix 가 있고, 현재 prefix 와 다르면 무시
            cached_prefix = data.get("app_key_prefix", "")
            if self._app_key_prefix and cached_prefix:
                if self._app_key_prefix != cached_prefix:
                    logger.warning("🔑 [%s] 앱키 변경 감지 → 기존 토큰 캐시 폐기", self._mode_str)
                    return

            if token and exp_dt:
                self._token  = token
                self._expiry = exp_dt
        except Exception as e:
            logger.debug("KIS 파일 캐시 로드 실패 [%s]: %s", self._mode_str, e)

    def _is_valid(self) -> bool:
        """
        API 사용 가능 여부. 장중 선제 발급 트리거용 10분 마진 없음.
        만료 SAFETY초 전까지만 유효 → 그 전에는 tokenP 안 침.
        """
        if not self._token or not self._expiry:
            return False
        return datetime.now() < self._expiry - timedelta(seconds=EXPIRE_SAFETY_SEC)

    # ── 공개 ──────────────────────────────────────────────────────
    def get_token(self) -> Optional[str]:
        """
        유효한 토큰 반환. 완전 만료(±60초) 시에만 비상 갱신.
        평소 발급은 장전 ensure_both_tokens / 타이머가 담당.

        흐름:
          1. 메모리 캐시 유효 → 즉시 반환 (파일 읽기 없음)
          2. 메모리 만료 → 파일 재로드 (다른 프로세스가 갱신했을 수 있음)
          3. 파일도 만료 → ensure_token() 으로 신규 발급 + 파일/메모리 갱신
          4. 발급 백오프 중이면 (EGW00133 직후 등) 만료된 토큰이라도 반환
             → 호출자가 None 받고 또 호출 → 또 None ... 무한루프 방지
        """
        with self._lock:
            if self._is_valid():
                return self._token
            # 다른 프로세스가 갱신했을 수 있으니 파일 재로드
            self._load_from_file()
            if self._is_valid():
                logger.info("🔑 %s 토큰 파일 갱신 감지 → 메모리 갱신", self._mode_str)
                return self._token

        # 백오프 중이면 발급 시도 자체를 안 함 (호출자 폭주 차단)
        allowed, wait_s = _can_issue_now(self._is_mock)
        if not allowed:
            logger.debug(
                "⏸ %s 토큰 발급 백오프 중 (%.1fs 남음) → 만료 캐시 토큰 임시 반환",
                self._mode_str, wait_s,
            )
            # 만료 토큰이라도 반환 (None 보다 안전 — 호출자 None 폭주 방지)
            return self._token

        # 파일도 만료 → 잠금 기반 신규 발급 (ensure 는 세션커버 판정)
        ok = ensure_token(self._is_mock)
        if ok:
            with self._lock:
                self._load_from_file()
            logger.info(
                "🔄 %s 토큰 비상 갱신 완료 (만료/세션미커버) → 앞8자: %s…",
                self._mode_str,
                (self._token[:8] if self._token else "N/A"),
            )
            return self._token
        # 발급 실패 — 백오프는 _issue_token 안에서 이미 설정됨.
        # 만료된 메모리 토큰이라도 있으면 반환 (호출자가 None 보고 또 두드리는 것 방지).
        return self._token

    def status(self) -> dict:
        """현재 상태 반환 (get_token_status 래퍼, 남은 시간 포함)"""
        return get_token_status(self._is_mock)


# ─────────────────────────────────────────────────────────────────────────────
# CLI: 상태 확인 / 강제 갱신
# ─────────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import argparse
    import logging as _lg

    _lg.basicConfig(
        level=_lg.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%H:%M:%S",
    )

    parser = argparse.ArgumentParser(description="KIS 토큰 관리자")
    parser.add_argument(
        "--refresh",
        action="store_true",
        help="캐시 삭제 후 강제 갱신 (1일1회 한도 주의 — 평소 사용 금지)",
    )
    args = parser.parse_args()

    _dl = _session_cover_deadline()
    print("\n" + "=" * 55)
    print("  KIS 토큰 상태")
    print(f"  세션커버 기준: 당일 {_dl.strftime('%H:%M')} 까지 유효해야 재사용")
    print("=" * 55)
    for is_mock, label in [(False, "실전"), (True, "모의")]:
        s = get_token_status(is_mock)
        if s["valid"]:
            flag = "✅ 세션커버"
        elif s.get("usable"):
            flag = "⚠️ 사용가능·세션미커버"
        else:
            flag = "❌ 만료/없음"
        print(f"  {label:4s} | {flag} | 만료: {s['expires']} ({s['expires_in_h']:+.1f}h)")
    print("=" * 55)

    if args.refresh:
        print("\n강제 갱신 시도 중...")
        # 캐시 무효화 후 재발급
        for path in [CACHE_REAL, CACHE_MOCK]:
            if path.exists():
                path.unlink()
                print(f"  캐시 삭제: {path.name}")

    result = ensure_both_tokens()
    print("\n갱신 결과:", result, "\n")