auto-stock/backend/app/services/tenant/registry.py

"""사용자 실행 컨텍스트 레지스트리 (설계 문서 §2, §5.6).

`UserContext` 조립(자격증명 복호화 → KIS 클라이언트 풀 → 리스크 매니저 팩토리 →
상태 복원)을 한곳에 모아 캐싱한다. FastAPI 의존성 `get_user_context()`도 여기서
제공한다(트레이딩 실행/엔드포인트 양쪽에서 재사용).

- DB 세션은 항상 호출 측이 넘겨준다 — 이 모듈은 `AsyncSessionLocal()`을 직접
  열지 않는다(§CLAUDE.md 아키텍처 원칙 1). `iter_active()`만 예외적으로 스케줄러가
  세션을 만들어 넘겨주는 것을 전제로 `db` 인자를 필수로 받는다.
- 복호화 실패 등 개별 사용자 오류는 그 사용자만 None/스킵 처리하고 다른 사용자에
  전파하지 않는다.
"""
import logging
import time

import asyncio

from fastapi import Depends, HTTPException
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.config import settings
from app.core.crypto import CredentialDecryptError, decrypt
from app.core.database import get_db
from app.core.tenant_auth import CurrentUser, get_current_user, get_legacy_owner
from app.models.models import UserKisCredential
from app.services.tenant.context import UserContext
from app.services.tenant.kis_pool import kis_client_pool
from app.services.tenant.risk_factory import (
    create_loss_tracker_for_user,
    create_mdd_manager_for_user,
    load_user_risk_state,
)

logger = logging.getLogger(__name__)

_LEGACY_OWNER_TTL_SEC = 60
_ACTIVE_LIST_TTL_SEC = 60


class UserContextRegistry:
    """사용자별 `UserContext` 캐시 + 조립 로직."""

    def __init__(self) -> None:
        self._contexts: dict[int, UserContext] = {}
        self._lock = asyncio.Lock()

        self._legacy_ctx: UserContext | None = None
        self._legacy_owner_id: int | None = None
        self._legacy_checked_at: float = 0.0
        self._legacy_lock = asyncio.Lock()
        self._legacy_owner_confirmed: bool = False

        self._active_cache: list[UserContext] | None = None
        self._active_cached_at: float = 0.0

    # ── 사용자별 컨텍스트 조립 ────────────────────────────────────────────
    async def get(self, user_id: int, db: AsyncSession) -> UserContext | None:
        """캐시 히트면 즉시 반환, 미스면 DB에서 자격증명을 읽어 새로 조립한다.

        M4: `asyncio.Lock` 은 캐시 dict 갱신 구간(더블체크 패턴)만 감싼다 — DB I/O,
        PBKDF2 복호화(스레드로 위임), KIS 클라이언트 생성, Redis/DB 리스크 상태 복원은
        모두 락 밖에서 수행해 이벤트 루프 직렬화를 피한다. 동시에 같은 user_id 로 호출되면
        중복 조립이 일어날 수 있으나(락 밖 조립) 마지막에 dict 갱신만 락으로 보호하고
        먼저 캐시된 인스턴스를 우선한다.
        """
        cached = self._contexts.get(user_id)
        if cached is not None:
            return cached

        result = await db.execute(
            select(UserKisCredential).where(UserKisCredential.user_id == user_id)
        )
        cred = result.scalar_one_or_none()
        if cred is None or not cred.is_enabled:
            logger.debug(f"[UserContextRegistry] user={user_id} 자격증명 없음/비활성 — None 반환")
            return None

        try:
            app_key = await asyncio.to_thread(decrypt, cred.app_key_enc, key_version=cred.key_version)
            app_secret = await asyncio.to_thread(decrypt, cred.app_secret_enc, key_version=cred.key_version)
            account_no = await asyncio.to_thread(decrypt, cred.account_no_enc, key_version=cred.key_version)
        except CredentialDecryptError as e:
            logger.error(f"[UserContextRegistry] user={user_id} 자격증명 복호화 실패: {e}")
            raise HTTPException(
                status_code=400, detail="자격증명 복호화 실패 — 서버 키 변경 여부 확인",
            ) from e

        try:
            kis = await kis_client_pool.get_or_create(
                user_id=user_id,
                app_key=app_key,
                app_secret=app_secret,
                account_no=account_no,
                mode=cred.mode,
            )
        except Exception as e:
            logger.error(f"[UserContextRegistry] user={user_id} KIS 클라이언트 생성 실패: {e}")
            return None

        mdd = create_mdd_manager_for_user(user_id)
        loss_tracker = create_loss_tracker_for_user(user_id)
        await load_user_risk_state(user_id, mdd, loss_tracker, db)

        ctx = UserContext(
            user_id=user_id,
            username="",
            mode=cred.mode,
            kis=kis,
            mdd=mdd,
            loss_tracker=loss_tracker,
            auto_trade_enabled=bool(cred.auto_trade_enabled),
            is_legacy=False,
        )
        async with self._lock:
            existing = self._contexts.get(user_id)
            if existing is not None:
                return existing
            self._contexts[user_id] = ctx
            return ctx

    # ── legacy(단일 계좌) 컨텍스트 ────────────────────────────────────────
    async def legacy(self, db: AsyncSession | None = None) -> UserContext:
        """기존 전역 싱글톤(`korea_investment_api`/`mdd_manager`/`loss_tracker`)을 감싼 컨텍스트.

        owner id는 60초 TTL로 재확인한다(관리자 계정이 나중에 생기는 초기 부팅 구간 대응).
        owner 가 아직 없으면(관리자 미존재) `user_id=None`으로 반환하고 TTL 캐시를 걸지
        않는다(M1) — `users.id=0`은 존재하지 않는 FK 위반 값이라 0 대신 None 으로 통일한다.
        """
        now = time.monotonic()
        if (
            self._legacy_ctx is not None
            and self._legacy_owner_confirmed
            and (now - self._legacy_checked_at) < _LEGACY_OWNER_TTL_SEC
        ):
            return self._legacy_ctx

        async with self._legacy_lock:
            now = time.monotonic()
            if (
                self._legacy_ctx is not None
                and self._legacy_owner_confirmed
                and (now - self._legacy_checked_at) < _LEGACY_OWNER_TTL_SEC
            ):
                return self._legacy_ctx

            owner_id: int | None = None
            if db is not None:
                try:
                    owner = await get_legacy_owner(db)
                    if owner is not None:
                        owner_id = owner.id
                except Exception as e:
                    logger.warning(f"[UserContextRegistry] legacy owner 조회 실패: {e}")

            from app.services.korea_investment_api import korea_investment_api
            from app.services.mdd_manager import mdd_manager
            from app.services.loss_tracker import loss_tracker

            self._legacy_ctx = UserContext(
                user_id=owner_id,
                username="legacy",
                mode=settings.KIS_MODE,
                kis=korea_investment_api,
                mdd=mdd_manager,
                loss_tracker=loss_tracker,
                auto_trade_enabled=settings.AUTO_TRADE_ENABLED,
                is_legacy=True,
            )
            self._legacy_owner_id = owner_id
            self._legacy_owner_confirmed = owner_id is not None
            self._legacy_checked_at = now
            return self._legacy_ctx

    # ── 인증 사용자 → 컨텍스트 해석 ───────────────────────────────────────
    async def resolve(self, current_user: CurrentUser, db: AsyncSession) -> UserContext:
        """`CurrentUser` → `UserContext`. 실패 시 400 HTTPException."""
        legacy_ctx = await self.legacy(db)

        use_legacy = (
            not settings.MULTI_TENANT_ENABLED
            or current_user.is_api_key
            or (current_user.id is not None and current_user.id == legacy_ctx.user_id)
        )
        if use_legacy:
            return legacy_ctx

        if current_user.id is None:
            raise HTTPException(status_code=400, detail="KIS 계좌가 등록·검증되지 않았습니다. 내 계좌 설정에서 등록하세요")

        ctx = await self.get(current_user.id, db)
        if ctx is None:
            raise HTTPException(status_code=400, detail="KIS 계좌가 등록·검증되지 않았습니다. 내 계좌 설정에서 등록하세요")
        return ctx

    # ── 2단계 스케줄러용: 활성 사용자 전체 순회 ───────────────────────────
    async def iter_active(self, db: AsyncSession) -> list[UserContext]:
        """`is_enabled and auto_trade_enabled` 인 모든 사용자의 컨텍스트 (60초 TTL 캐시).

        지금은 조립 로직만 제공 — 실제 2단계 스케줄러 연결은 이번 범위 밖.
        """
        now = time.monotonic()
        if self._active_cache is not None and (now - self._active_cached_at) < _ACTIVE_LIST_TTL_SEC:
            return self._active_cache

        result = await db.execute(
            select(UserKisCredential).where(
                UserKisCredential.is_enabled.is_(True),
                UserKisCredential.auto_trade_enabled.is_(True),
            )
        )
        creds = result.scalars().all()

        contexts: list[UserContext] = []
        for cred in creds:
            try:
                ctx = await self.get(cred.user_id, db)
            except Exception as e:
                logger.error(f"[UserContextRegistry] iter_active: user={cred.user_id} 조립 실패: {e}")
                ctx = None
            if ctx is not None:
                contexts.append(ctx)

        self._active_cache = contexts
        self._active_cached_at = now
        return contexts

    # ── 캐시 무효화 ───────────────────────────────────────────────────────
    async def invalidate(self, user_id: int) -> None:
        """자격증명 변경/삭제 시 credentials API가 호출 — 컨텍스트 캐시 + KIS 풀 정리."""
        async with self._lock:
            self._contexts.pop(user_id, None)
        await kis_client_pool.invalidate(user_id)
        self._active_cache = None

    async def invalidate_all(self) -> None:
        async with self._lock:
            self._contexts.clear()
        self._active_cache = None


# 모듈 싱글톤
user_context_registry = UserContextRegistry()


async def get_user_context(
    user: CurrentUser = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
) -> UserContext:
    """FastAPI 의존성 — 세션/API 키 인증 사용자를 `UserContext`로 해석."""
    return await user_context_registry.resolve(user, db)