auto-stock/backend/app/services/loss_tracker.py

"""
연속 손실 추적 - CLAUDE.md 섹션 7.4

동일 종목에서 3회 이상 연속 손실 시 30거래일 동안 신규 진입 금지.
재시작 후에도 ban 상태 유지: DB(stock_bans 테이블) 영속화.
"""
import asyncio
import logging
from datetime import datetime, timedelta, timezone
from typing import Dict, Optional

logger = logging.getLogger(__name__)

# tz-aware/naive 경계 정규화 — app/core/dt_utils 참고 (db_worker.py와 공유)
from app.core.dt_utils import to_naive_utc as _to_naive_utc, to_aware_utc as _to_aware_utc  # noqa: E402


class ConsecutiveLossTracker:
    """종목별 연속 손실 추적 및 쿨링오프 관리"""

    # 멀티테넌트: risk_factory.create_loss_tracker_for_user()가 인스턴스 생성 후
    # 세팅하는 소유자 user_id. 레거시 싱글톤은 None으로 유지(전역 stock_bans 사용).
    user_id: Optional[int] = None

    def __init__(self, ban_count: int = 3, ban_calendar_days: int = 42):
        """
        Args:
            ban_count: 연속 손실 횟수 임계값 (기본 3)
            ban_calendar_days: 금지 기간 (30거래일 ≈ 42캘린더일)
        """
        self.ban_count = ban_count
        self.ban_calendar_days = ban_calendar_days
        # {stock_code: {"consecutive_losses": int, "ban_until": datetime|None}}
        self._tracker: Dict[str, dict] = {}
        self._lock = asyncio.Lock()  # [FIX] 동시 접근 방지용 비동기 락

    def record_trade_result(self, stock_code: str, is_loss: bool):
        """매매 결과 인메모리 기록 (DB 저장은 save_to_db 별도 호출)

        Args:
            stock_code: 종목코드
            is_loss: 손실 여부
        """
        if stock_code not in self._tracker:
            self._tracker[stock_code] = {"consecutive_losses": 0, "ban_until": None}

        entry = self._tracker[stock_code]
        if is_loss:
            entry["consecutive_losses"] += 1
            if entry["consecutive_losses"] >= self.ban_count:
                entry["ban_until"] = datetime.now(timezone.utc) + timedelta(days=self.ban_calendar_days)
                logger.warning(
                    f"{stock_code}: {entry['consecutive_losses']}회 연속 손실 → "
                    f"{entry['ban_until'].strftime('%Y-%m-%d')}까지 진입 금지"
                )
        else:
            entry["consecutive_losses"] = 0
            entry["ban_until"] = None

    async def save_to_db(self, db=None, stock_code: str = None):
        """인메모리 상태를 DB에 영속화 (record_trade_result 후 호출)

        NOTE: db 파라미터는 하위호환성을 위해 유지되지만 사용되지 않습니다.
        실제 DB 저장은 db_worker.py의 db_task_queue를 통해 처리됩니다.
        """
        if stock_code is None:
            return

        try:
            from app.services.db_worker import enqueue_stock_ban_upsert

            entry = self._tracker.get(stock_code)
            if entry is None:
                return

            # 큐에 UPSERT 작업 등록 (db_worker가 비동기적으로 처리)
            # ban_until 컬럼은 tz-naive이므로 저장 직전 naive-UTC로 정규화한다.
            await enqueue_stock_ban_upsert(
                stock_code=stock_code,
                consecutive_losses=entry["consecutive_losses"],
                ban_until=_to_naive_utc(entry["ban_until"]),
            )
            logger.debug(f"loss_tracker DB 저장 큐 등록: {stock_code}")
        except Exception as e:
            logger.error(f"loss_tracker DB 저장 큐 등록 실패 ({stock_code}): {e}")

    async def save_to_db_v2(self, user_id: int, stock_code: str = None):
        """[v2] 사용자 스코프 인메모리 상태를 stock_bans_v2 에 영속화.

        save_to_db()와 동일 패턴이지만 db_worker의 upsert_stock_ban_v2 task로
        (user_id, stock_code) 단위 UPSERT한다. 기존 save_to_db()는 불변.
        """
        if stock_code is None:
            return

        try:
            from app.services.db_worker import enqueue_stock_ban_upsert_v2

            entry = self._tracker.get(stock_code)
            if entry is None:
                return

            # stock_bans_v2.ban_until 도 tz-naive 컬럼 — 저장 직전 naive-UTC로 정규화.
            await enqueue_stock_ban_upsert_v2(
                user_id=user_id,
                stock_code=stock_code,
                consecutive_losses=entry["consecutive_losses"],
                ban_until=_to_naive_utc(entry["ban_until"]),
            )
            logger.debug(f"loss_tracker_v2 DB 저장 큐 등록: user={user_id} {stock_code}")
        except Exception as e:
            logger.error(f"loss_tracker_v2 DB 저장 큐 등록 실패 (user={user_id}, {stock_code}): {e}")

    async def load_from_db_v2(self, db, user_id: int):
        """[v2] 앱 시작/컨텍스트 생성 시 stock_bans_v2 에서 사용자 스코프 ban 상태 복원."""
        try:
            from sqlalchemy import select
            from app.models.models import StockBanV2

            result = await db.execute(select(StockBanV2).where(StockBanV2.user_id == user_id))
            records = result.scalars().all()

            now = datetime.now(timezone.utc)
            loaded = 0
            for rec in records:
                # DB의 ban_until은 tz-naive → tz-aware now와 비교하기 전 정규화.
                ban_until = _to_aware_utc(rec.ban_until)
                if ban_until and now >= ban_until:
                    ban_until = None
                    consecutive_losses = 0
                else:
                    consecutive_losses = rec.consecutive_losses

                self._tracker[rec.stock_code] = {
                    "consecutive_losses": consecutive_losses,
                    "ban_until": ban_until,
                }
                if ban_until:
                    loaded += 1
                    logger.info(
                        f"ban_v2 복원: user={user_id} {rec.stock_code} → "
                        f"{ban_until.strftime('%Y-%m-%d')}까지 진입 금지"
                    )

            logger.info(
                f"loss_tracker_v2 DB 로드 완료: user={user_id} {len(records)}건 중 {loaded}건 활성 ban"
            )
        except Exception:
            logger.exception(f"loss_tracker_v2 DB 로드 실패 (user={user_id})")

    async def load_from_db(self, db):
        """앱 시작 시 DB에서 ban 상태 복원"""
        try:
            from sqlalchemy import select
            from app.models.models import StockBan

            result = await db.execute(select(StockBan))
            records = result.scalars().all()

            now = datetime.now(timezone.utc)
            loaded = 0
            for rec in records:
                # 만료된 ban은 초기화. DB의 ban_until은 tz-naive → 비교 전 정규화.
                ban_until = _to_aware_utc(rec.ban_until)
                if ban_until and now >= ban_until:
                    ban_until = None
                    consecutive_losses = 0
                else:
                    consecutive_losses = rec.consecutive_losses

                self._tracker[rec.stock_code] = {
                    "consecutive_losses": consecutive_losses,
                    "ban_until": ban_until,
                }
                if ban_until:
                    loaded += 1
                    logger.info(
                        f"ban 복원: {rec.stock_code} → {ban_until.strftime('%Y-%m-%d')}까지 진입 금지"
                    )

            logger.info(f"loss_tracker DB 로드 완료: {len(records)}건 중 {loaded}건 활성 ban")
        except Exception:
            logger.exception("loss_tracker DB 로드 실패")

    def is_banned(self, stock_code: str) -> bool:
        """해당 종목이 진입 금지 상태인지 확인"""
        entry = self._tracker.get(stock_code)
        if not entry or not entry["ban_until"]:
            return False
        if datetime.now(timezone.utc) >= entry["ban_until"]:
            # 금지 기간 만료 → 리셋
            entry["consecutive_losses"] = 0
            entry["ban_until"] = None
            return False
        return True

    async def is_banned_safe(self, stock_code: str) -> bool:
        """[v2] 비동기 락 보호 is_banned — 동시 접근이 우려되는 경로에서 사용"""
        async with self._lock:
            return self.is_banned(stock_code)

    async def record_trade_result_safe(self, stock_code: str, is_loss: bool):
        """[v2] 비동기 락 보호 record_trade_result"""
        async with self._lock:
            self.record_trade_result(stock_code, is_loss)

    async def get_all_bans_safe(self) -> list:
        """[v2] Lock 보호 하에 현재 활성 ban 종목 코드 목록 반환.

        Returns:
            ban_until이 아직 유효한 종목코드 리스트
        """
        async with self._lock:
            now = datetime.now(timezone.utc)
            return [
                code for code, entry in self._tracker.items()
                if entry.get("ban_until") and now < entry["ban_until"]
            ]

    def get_ban_info(self, stock_code: str) -> dict:
        """종목의 연속 손실 및 금지 정보 조회"""
        entry = self._tracker.get(stock_code, {"consecutive_losses": 0, "ban_until": None})
        return {
            "stock_code": stock_code,
            "consecutive_losses": entry["consecutive_losses"],
            "is_banned": self.is_banned(stock_code),
            "ban_until": entry["ban_until"].isoformat() if entry["ban_until"] else None,
        }


# 싱글톤 인스턴스 — settings에서 ban_count 주입
def _create_loss_tracker() -> ConsecutiveLossTracker:
    try:
        from app.core.config import settings
        return ConsecutiveLossTracker(
            ban_count=settings.CONSECUTIVE_LOSS_BAN_COUNT,
            ban_calendar_days=42,  # 30거래일 ≈ 42캘린더일
        )
    except Exception:
        return ConsecutiveLossTracker()


loss_tracker = _create_loss_tracker()