auto-stock/backend/app/main.py

"""
FastAPI 애플리케이션 진입점

- CORS, 미들웨어 설정
- WebSocket 실시간 시세 브로드캐스트
- 라우터 등록 및 startup/shutdown 이벤트 관리
"""
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
import contextlib
import asyncio
import json
import logging
import time
import warnings

warnings.filterwarnings("ignore", message=".*sklearn.utils.parallel.delayed.*")
import pytz

from fastapi import Depends
from app.core.config import settings
from app.core.security import require_api_key as _require_admin_key
from app.api import api_router
from app.services.websocket_manager import websocket_manager
from app.core.json_helpers import _dumps
from app.services.db_worker import db_writer_worker
from app.core.database import engine, Base, AsyncSessionLocal
from app.models import models  # 모델 임포트
from app.core.market_schedule import is_market_open

# 로깅 설정
import os
import logging.handlers
from app.core.logging_config import KSTFormatter
_log_level = getattr(logging, settings.LOG_LEVEL.upper(), logging.INFO)
_log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
_log_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "logs")
os.makedirs(_log_dir, exist_ok=True)
_log_file = os.path.join(_log_dir, "backend.log")

_formatter = KSTFormatter(_log_format)
_handlers: list[logging.Handler] = [
    logging.StreamHandler(),
    logging.handlers.RotatingFileHandler(
        _log_file, maxBytes=10 * 1024 * 1024, backupCount=5, encoding="utf-8"
    ),
]
for _h in _handlers:
    _h.setFormatter(_formatter)

logging.basicConfig(level=_log_level, handlers=_handlers)

logger = logging.getLogger(__name__)

# Import scheduler functions from dedicated module
from app.services.schedulers import (
    batch_analysis_scheduler,
    post_market_analyzer_scheduler,
    krx_data_collector_scheduler,
    news_collect_scheduler,
    pre_market_us_scheduler,
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """애플리케이션 시작/종료 시 실행될 코드"""
    logger.info("Starting Auto Stock Trading System...")
    logger.info(f"KIS API Mode: {'모의투자 (VIRTUAL)' if settings.is_virtual_trading else '⚠️ 실전투자 (REAL)'}")
    logger.info(f"KIS Base URL: {settings.KOREA_INVESTMENT_API_BASE_URL}")
    _acct = settings.KOREA_INVESTMENT_ACCOUNT_NUMBER
    _masked_acct = f"{_acct[:4]}****{_acct[-2:]}" if len(_acct) >= 6 else "***"
    logger.info(f"KIS Account: {_masked_acct}")

    # ── 필수 환경변수 검증 (조기 경보) ──────────────────────────────────────────
    for warn_msg in settings.validate_required():
        logger.warning(f"⚠️ [Config] {warn_msg}")

    # DB 테이블 자동 생성
    # 컬럼 추가/테이블 생성 DDL 은 Alembic migration 으로 관리됨.
    # 배포 전 `cd backend && alembic upgrade head` 를 실행할 것.
    # (alembic/versions/0001_initial_schema_columns.py)
    try:
        logger.info(f"Connecting to database: {settings.DATABASE_URL.split('@')[1] if '@' in settings.DATABASE_URL else 'unknown'}")
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        logger.info("✅ Database tables created successfully")

        try:
            from app.services.loss_tracker import loss_tracker
            async with AsyncSessionLocal() as session:
                await loss_tracker.load_from_db(session)
        except Exception as e:
            logger.warning(f"⚠️ loss_tracker 초기화 실패 (non-critical): {e}")

        # ── 전략 파라미터 DB 초기화 (config.py 기본값 → strategy_configs 테이블 seed) ──
        try:
            from app.services.strategy_config_service import strategy_cfg
            await strategy_cfg.auto_init()
        except Exception as e:
            logger.warning(f"⚠️ strategy_cfg 초기화 실패 (non-critical): {e}")
        # ────────────────────────────────────────────────────────────────────────────

        try:
            from app.models.models import StockAnalysis, TradeHistory, Portfolio
            from sqlalchemy import delete
            from datetime import datetime, timedelta

            async with AsyncSessionLocal() as session:
                _kst = pytz.timezone("Asia/Seoul")
                _today_kst = datetime.now(_kst).date()
                cutoff_date = _today_kst - timedelta(days=7)
                result = await session.execute(
                    delete(StockAnalysis).where(StockAnalysis.analysis_date < cutoff_date)
                )
                if result.rowcount > 0:
                    logger.info(f"🗑️ Deleted {result.rowcount} old stock_analysis records")

                cutoff_datetime = _today_kst - timedelta(days=1)
                result = await session.execute(
                    delete(TradeHistory).where(
                        TradeHistory.status.in_(['PENDING', 'FAILED']),
                        TradeHistory.created_at < cutoff_datetime,
                    )
                )
                if result.rowcount > 0:
                    logger.info(f"🗑️ Deleted {result.rowcount} old PENDING/FAILED trades")

                result = await session.execute(
                    delete(Portfolio).where(Portfolio.quantity <= 0)
                )
                if result.rowcount > 0:
                    logger.info(f"🗑️ Deleted {result.rowcount} empty portfolio records")

                await session.commit()
                logger.info("✅ Test data cleanup completed")
        except Exception as e:
            logger.warning(f"⚠️ Test data cleanup failed (non-critical): {e}")

    except ConnectionRefusedError as e:
        logger.error(f"❌ Database connection refused: {e}")
        raise
    except Exception as e:
        logger.error(f"❌ Failed to create database tables: {e}")
        import traceback
        logger.error(traceback.format_exc())
        raise

    # ── Alembic 자동 마이그레이션 (서버 이전 시 DB 스키마 자동 동기화) ──────
    try:
        from alembic.config import Config as AlembicConfig
        from alembic import command as alembic_command
        import os
        alembic_ini = os.path.join(os.path.dirname(os.path.dirname(__file__)), "alembic.ini")
        if os.path.exists(alembic_ini):
            alembic_cfg = AlembicConfig(alembic_ini)
            alembic_cfg.set_main_option("script_location", os.path.join(os.path.dirname(os.path.dirname(__file__)), "alembic"))
            alembic_command.upgrade(alembic_cfg, "head")
            logger.info("✅ Alembic migration completed (upgrade head)")
    except Exception as e:
        logger.warning(f"⚠️ Alembic migration 실패 (non-critical): {e}")

    # ── KRX 일봉 데이터 초기 수집 (DB 비어있을 때만 백그라운드 실행) ──────
    try:
        from app.models.models import StockDailyPrice
        from sqlalchemy import select as _select, func as _func
        async with AsyncSessionLocal() as _check_session:
            _count = await _check_session.execute(
                _select(_func.count()).select_from(StockDailyPrice)
            )
            _total = _count.scalar() or 0
        if _total == 0 and settings.DATAGOKR_API_KEY:
            async def _initial_krx_collect():
                await asyncio.sleep(10)  # 부팅 안정화 대기
                try:
                    from collect_krx_data import collect_all
                    logger.info("📥 [KRX] 초기 데이터 수집 시작 (stock_daily_prices 비어있음)...")
                    await collect_all(update_mode=True)
                    logger.info("✅ [KRX] 초기 데이터 수집 완료")
                except Exception as e:
                    logger.warning(f"⚠️ [KRX] 초기 데이터 수집 실패 (non-critical): {e}")
            asyncio.create_task(_initial_krx_collect())
            logger.info("📥 [KRX] stock_daily_prices 비어있음 — 백그라운드 초기 수집 예약됨")
        elif _total == 0:
            logger.warning("⚠️ [KRX] stock_daily_prices 비어있으나 DATAGOKR_API_KEY 미설정 — 네이버 차트 폴백 사용")
        else:
            logger.info(f"📊 [KRX] stock_daily_prices: {_total:,}행 확인")
    except Exception as e:
        logger.warning(f"⚠️ [KRX] 초기 데이터 확인 실패 (non-critical): {e}")

    try:
        from app.api.endpoints.day_trading import restore_positions_from_db
        await restore_positions_from_db()
    except Exception as e:
        logger.warning(f"⚠️ 데이트레이딩 포지션 복구 실패 (non-critical): {e}")

    # ── httpx 싱글톤 클라이언트 초기화 ──────────────────────────────────────
    from app.api.endpoints.mcp import init_http_client, close_http_client
    await init_http_client()

    # ── Redis 클라이언트 초기화 ───────────────────────────────────────────────
    from app.core.redis import init_redis, close_redis
    await init_redis()

    # ── MDD 상태 복원 (Redis 영속화) ────────────────────────────────────────
    try:
        from app.services.mdd_manager import mdd_manager
        restored = await mdd_manager.load_from_redis()
        if not restored:
            logger.info("ℹ️  MDD Redis 상태 없음 — 초기 상태로 시작")
    except Exception as e:
        logger.warning(f"⚠️ MDD Redis 복원 실패 (non-critical): {e}")

    # 워커 + 스케줄러 시작 (C5: 자동 재시작 콜백 포함)
    def _restart_db_worker(task: asyncio.Task):
        if task.cancelled():
            return
        exc = task.exception()
        if exc:
            logger.error(f"db_worker 크래시: {exc}, 자동 재시작...")
            # Telegram 시스템 경고
            try:
                from app.services.telegram_notifier import telegram_notifier
                asyncio.ensure_future(telegram_notifier.notify_system_warning(
                    category="DB_WORKER",
                    message=f"db_worker 크래시 후 자동 재시작: {type(exc).__name__}: {exc}",
                ))
            except Exception:
                pass
            new_task = asyncio.create_task(db_writer_worker())
            new_task.add_done_callback(_restart_db_worker)
            app.state.db_worker_task = new_task

    app.state.db_worker_task = asyncio.create_task(db_writer_worker())
    app.state.db_worker_task.add_done_callback(_restart_db_worker)

    # ── 스윙 / 데이트레이딩 분리 스케줄러 (신규 아키텍처) ────────────────────
    from app.services.swing.scheduler import swing_scheduler
    from app.services.daytrading.scheduler import daytrading_scheduler
    app.state.swing_trade_task = asyncio.create_task(swing_scheduler())
    app.state.day_trade_task = asyncio.create_task(daytrading_scheduler())

    # 배치 분석 / 일봉 수집 / 장 후 분석 (기존 유지)
    app.state.batch_analysis_task = asyncio.create_task(batch_analysis_scheduler())
    app.state.krx_collector_task = asyncio.create_task(krx_data_collector_scheduler())
    app.state.post_market_task = asyncio.create_task(post_market_analyzer_scheduler())

    # ── 미국장 자동매매 스케줄러 (US_TRADING_ENABLED=True 시 활성화) ──────────
    if getattr(settings, 'US_TRADING_ENABLED', False):
        from app.services.us_market_scheduler import us_auto_trade_scheduler
        app.state.us_trade_task = asyncio.create_task(us_auto_trade_scheduler())
        logger.info("🇺🇸 US market auto-trade scheduler started (US_TRADING_ENABLED=True)")
    else:
        logger.info("⏭️  US market auto-trade disabled (US_TRADING_ENABLED=False in .env)")
        app.state.us_trade_task = None

    # ── 뉴스 수집 스케줄러 (NEWS_ENABLED=True 시 활성화) ─────────────────────
    app.state.news_collect_task = asyncio.create_task(news_collect_scheduler())
    logger.info("News collect scheduler started")

    # ── 미장 프리마켓 심리 분석 (07:00 KST, 1회/일) ────────────────────────
    app.state.pre_market_us_task = asyncio.create_task(pre_market_us_scheduler())
    logger.info("🌍 Pre-market US sentiment scheduler registered (daily 07:00 KST)")

    # ── ML 재학습 스케줄러 ───────────────────────────────────────────────────
    from app.services.ml_retrain_pipeline import ml_retrain_scheduler
    app.state.ml_retrain_task = asyncio.create_task(ml_retrain_scheduler())
    logger.info("ML retrain scheduler started")

    # ── AI 전략 스케줄러 (4단계: 장 전/장중/매매직전/장 후) ──────────────────
    async def _ai_strategy_scheduler():
        """AI 4단계 전략 스케줄러.

        08:30  장 전 Sonnet 배치 분석 (스윙 50종목)
        08:50  데이트레이딩 종목 자동 선정 (Haiku)
        10분   장중 뉴스 모니터링 (Haiku, 악재 시 Sonnet)
        15:40  장 후 리뷰 + Telegram (Sonnet)
        """
        import pytz as _tz
        from datetime import datetime as _dt

        await asyncio.sleep(5)
        _kst = _tz.timezone("Asia/Seoul")
        _pre_market_done_today = None
        _daytrading_done_today = None
        _post_market_done_today = None
        _last_news_check = None

        logger.info("🤖 [AI Strategy] Scheduler started")

        while True:
            try:
                now = _dt.now(_kst)
                today = now.date()
                hm = now.hour * 100 + now.minute  # HHMM 형식

                # 평일만 (월~금)
                if now.weekday() >= 5:
                    await asyncio.sleep(300)
                    continue

                # 08:30 — 장 전 스윙 배치 분석
                if 830 <= hm < 845 and _pre_market_done_today != today:
                    try:
                        from app.services.ai_strategy_service import pre_market_analysis
                        result = await pre_market_analysis()
                        suitable = [s for s in result.get("suitable_stocks", []) if s.get("suitable")]
                        logger.info(f"🤖 [AI Strategy] 장 전 분석 완료: {len(suitable)}종목 적합")
                        _pre_market_done_today = today
                    except Exception as e:
                        logger.warning(f"[AI Strategy] 장 전 분석 실패: {e}")

                # 08:50 — 데이트레이딩 종목 선정
                if 850 <= hm < 900 and _daytrading_done_today != today:
                    if settings.DAY_TRADE_ENABLED:
                        try:
                            from app.services.ai_strategy_service import select_daytrading_candidates
                            candidates = await select_daytrading_candidates()
                            logger.info(f"🤖 [AI Strategy] 데이트레이딩 후보: {len(candidates)}종목")
                            _daytrading_done_today = today
                        except Exception as e:
                            logger.warning(f"[AI Strategy] 데이트레이딩 선정 실패: {e}")

                # 09:10~15:20 — 장중 뉴스 모니터링 (10분 주기)
                if 910 <= hm <= 1520:
                    should_check = (
                        _last_news_check is None
                        or (now - _last_news_check).total_seconds() >= 600
                    )
                    if should_check:
                        try:
                            from app.services.ai_strategy_service import monitor_news
                            alerts = await monitor_news()
                            if alerts:
                                logger.warning(f"🚨 [AI Strategy] 뉴스 악재 감지: {len(alerts)}건")
                            _last_news_check = now
                        except Exception as e:
                            logger.debug(f"[AI Strategy] 뉴스 모니터링 실패: {e}")

                # 15:40 — 장 후 리뷰
                if 1540 <= hm < 1600 and _post_market_done_today != today:
                    try:
                        from app.services.ai_strategy_service import post_market_review
                        review = await post_market_review()
                        grade = review.get("overall_grade", "?")
                        logger.info(f"🤖 [AI Strategy] 장 후 리뷰 완료: 등급 {grade}")
                        _post_market_done_today = today
                    except Exception as e:
                        logger.warning(f"[AI Strategy] 장 후 리뷰 실패: {e}")

            except Exception as e:
                logger.error(f"[AI Strategy] 스케줄러 오류: {e}")

            await asyncio.sleep(30)

    app.state.ai_strategy_task = asyncio.create_task(_ai_strategy_scheduler())
    logger.info("🤖 AI strategy scheduler started (pre-market/news/review)")

    # ── 적응형 전략 엔진 초기화 + 주기적 갱신 ────────────────────────────────
    async def _adaptive_strategy_scheduler():
        """적응형 전략 파라미터를 2분 주기로 갱신."""
        from app.services.adaptive_strategy import refresh_adaptive_params
        await asyncio.sleep(10)  # 서버 기동 후 10초 대기
        while True:
            try:
                await refresh_adaptive_params()
            except Exception as e:
                logger.warning("[Adaptive] 스케줄러 오류: %s", e)
            await asyncio.sleep(120)  # 2분 주기
    app.state.adaptive_strategy_task = asyncio.create_task(_adaptive_strategy_scheduler())
    logger.info("Adaptive strategy engine scheduler started")

    # KIS 실시간 스트림 시작 (설정으로 제어 가능)
    if settings.KIS_WS_ENABLED:
        from app.services.kis_realtime_stream import run_kis_realtime_stream
        app.state.kis_stream_task = asyncio.create_task(run_kis_realtime_stream())
        logger.info("📡 KIS 실시간 스트림 태스크 시작 (KIS_WS_ENABLED=True)")
    else:
        logger.info("⏭️  KIS 실시간 스트림 비활성화됨 (KIS_WS_ENABLED=False in .env)")
        app.state.kis_stream_task = None

    yield

    # 태스크 정리
    for task_attr in (
        "swing_trade_task",
        "batch_analysis_task",
        "day_trade_task",
        "krx_collector_task",
        "post_market_task",
        "db_worker_task",
        "kis_stream_task",
        "us_trade_task",
        "news_collect_task",
        "ml_retrain_task",
        "pre_market_us_task",
        "adaptive_strategy_task",
        "ai_strategy_task",
    ):
        try:
            task = getattr(app.state, task_attr, None)
            if task:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task
        except Exception as e:
            logger.warning(f"Failed to stop {task_attr} gracefully: {e}")

    # ── httpx 클라이언트 종료 ─────────────────────────────────────────────────
    try:
        await close_http_client()
    except Exception as e:
        logger.warning(f"Failed to close http_client: {e}")

    # ── KIS API httpx 클라이언트 종료 ───────────────────────────────────────
    try:
        from app.services.korea_investment_api import korea_investment_api
        await korea_investment_api.close()
    except Exception as e:
        logger.warning(f"Failed to close KIS http_client: {e}")

    # ── Redis 클라이언트 종료 ─────────────────────────────────────────────────
    try:
        await close_redis()
    except Exception as e:
        logger.warning(f"Failed to close Redis: {e}")

    # ── Telegram 클라이언트 종료 ───────────────────────────────────────────────
    try:
        from app.services.telegram_notifier import telegram_notifier
        await telegram_notifier.close()
    except Exception as e:
        logger.warning(f"Failed to close Telegram notifier: {e}")

    logger.info("Shutting down Auto Stock Trading System...")
    await engine.dispose()


app = FastAPI(
    title=settings.APP_NAME,
    description="Claude AI와 MCP를 활용한 자동 주식 매매 시스템",
    version="1.0.0",
    lifespan=lifespan,
    debug=settings.DEBUG,
)

# CORS 설정
_cors_origins = getattr(settings, "CORS_ORIGINS", None)
if _cors_origins:
    _allowed_origins = [o.strip() for o in _cors_origins.split(",") if o.strip()]
else:
    _allowed_origins = [
        "http://localhost:5173",
        "http://localhost:3000",
        "http://127.0.0.1:5173",
        "http://127.0.0.1:3000",
    ]
app.add_middleware(
    CORSMiddleware,
    allow_origins=_allowed_origins,
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["Content-Type", "Authorization", "X-API-Key", "X-Requested-With"],
    expose_headers=["X-Request-Id"],
)
logger.info("CORS middleware configured")


# ── [C6] 간이 Rate Limiter (IP 기반, 인메모리) ──────────────────────────────
_rate_limit_store: dict[str, list[float]] = {}  # IP -> [timestamps]
_RATE_LIMIT_WINDOW = 60  # 초
_RATE_LIMIT_MAX_REQUESTS = 300  # 윈도우 내 최대 요청 수 (대시보드 폴링 고려)
_RATE_LIMIT_CLEANUP_INTERVAL = 100  # N 요청마다 stale IP 정리
_RATE_LIMIT_STALE_THRESHOLD = 3600  # 1시간 이상 미접근 IP 삭제
_rate_limit_counter = 0  # cleanup 트리거 카운터


@app.middleware("http")
async def rate_limit_middleware(request: Request, call_next):
    """IP 기반 간이 rate limiter — 인증 없는 엔드포인트 DoS 방어."""
    global _rate_limit_counter
    client_ip = request.client.host if request.client else "unknown"
    now = time.time()

    # /health, /ws/, OPTIONS(CORS preflight) 는 rate limit 제외
    path = request.url.path
    if path.startswith(("/health", "/ws/")) or request.method == "OPTIONS":
        return await call_next(request)

    timestamps = _rate_limit_store.get(client_ip, [])
    # 윈도우 밖 타임스탬프 제거
    timestamps = [t for t in timestamps if now - t < _RATE_LIMIT_WINDOW]
    timestamps.append(now)
    _rate_limit_store[client_ip] = timestamps

    if len(timestamps) > _RATE_LIMIT_MAX_REQUESTS:
        from fastapi.responses import JSONResponse
        return JSONResponse(
            status_code=429,
            content={"detail": "Too many requests. Please try again later."},
        )

    # 100 요청마다 stale IP 키 정리 (1시간 이상 미접근)
    _rate_limit_counter += 1
    if _rate_limit_counter >= _RATE_LIMIT_CLEANUP_INTERVAL:
        _rate_limit_counter = 0
        stale_ips = [
            ip for ip, ts_list in _rate_limit_store.items()
            if not ts_list or (now - max(ts_list)) > _RATE_LIMIT_STALE_THRESHOLD
        ]
        for ip in stale_ips:
            del _rate_limit_store[ip]
        if stale_ips:
            logger.debug(f"Rate limit store cleanup: {len(stale_ips)} stale IPs removed, {len(_rate_limit_store)} remaining")

    return await call_next(request)


_LOG_SKIP_PREFIXES = ("/health", "/ws/")


@app.middleware("http")
async def log_requests(request: Request, call_next):
    path = request.url.path
    # /health, /ws/ 는 폴링이 잦아 로그 노이즈가 크므로 debug 레벨로 낮춤
    if any(path.startswith(p) for p in _LOG_SKIP_PREFIXES):
        return await call_next(request)

    start_time = time.time()
    logger.info(f"🔵 Request: {request.method} {request.url}")
    response = await call_next(request)
    process_time = time.time() - start_time
    logger.info(
        f"🟢 Response: {request.method} {request.url} "
        f"- Status: {response.status_code} - Time: {process_time:.3f}s"
    )
    return response


app.include_router(api_router, prefix="/api/v1")


# ── 전역 예외 핸들러 ────────────────────────────────────────────────────────────
from fastapi import Request as _Request
from fastapi.responses import JSONResponse as _JSONResponse


@app.exception_handler(Exception)
async def global_exception_handler(_req: _Request, exc: Exception) -> _JSONResponse:
    """처리되지 않은 예외를 500 JSON으로 변환해 일관된 에러 응답을 보장한다."""
    logger.error(f"Unhandled exception: {_req.method} {_req.url} — {exc}", exc_info=True)
    content = {"detail": "내부 서버 오류가 발생했습니다. 잠시 후 다시 시도해 주세요."}
    # DEBUG 모드에서만 스택트레이스 노출 (운영 환경 정보 유출 방지)
    if settings.DEBUG:
        content["error"] = str(exc)
    return _JSONResponse(status_code=500, content=content)


@app.get("/")
async def root():
    return {"message": "Auto Stock Trading System API", "version": "1.0.0", "status": "running"}


@app.get("/health")
async def health_check():
    """DB / Redis 연결 상태를 포함한 헬스체크."""
    from sqlalchemy import text as _sa_text
    from app.core.redis import get_redis

    db_ok = False
    redis_ok = False

    try:
        async with engine.connect() as conn:
            await conn.execute(_sa_text("SELECT 1"))
        db_ok = True
    except Exception:
        pass

    try:
        r = get_redis()
        if r:
            await r.ping()
            redis_ok = True
    except Exception:
        pass

    status = "healthy" if (db_ok and redis_ok) else "degraded"
    return {
        "status": status,
        "db": "ok" if db_ok else "error",
        "redis": "ok" if redis_ok else "error",
    }


@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    """WebSocket 연결 엔드포인트.

    인증: 쿼리 파라미터 ?token= 또는 첫 메시지로 {"type": "auth", "token": "..."}
    개발 모드(DEBUG=true)에서는 인증 없이 접속 허용.

    클라이언트→서버 지원 프로토콜:
      {"type": "ping"}                                       → {"type": "pong"}
      {"type": "subscribe", "stock_codes": [...]}           → ACK
      {"type": "unsubscribe", "stock_codes": [...]}         → ACK
      {"type": "set_subscriptions", "stock_codes": [...]}   → ACK
      {"type": "get_chart", "stock_code": "...", "period": "D", "limit": 200}
                                                            → {"type": "chart_data", ...}
      그 외 메시지                                           → echo
    """
    # WebSocket 인증: 운영 환경에서는 토큰 검증 (API_KEY 또는 SECRET_KEY 허용)
    if not settings.DEBUG:
        import hmac
        from app.core.security import _get_api_key
        token = websocket.query_params.get("token", "")
        expected = _get_api_key()
        if not token or not hmac.compare_digest(token.encode(), expected.encode()):
            await websocket.close(code=4001, reason="Unauthorized")
            return

    await websocket_manager.connect(websocket, client_id)

    # 연결 즉시 보유종목 자동 구독 (set_subscriptions 누락 대비 폴백)
    async def _auto_subscribe_portfolio(cid: str):
        try:
            from app.api.endpoints.portfolio import portfolio_cache
            holdings = portfolio_cache.holdings
            if holdings:
                codes = [h.stock_code for h in holdings if h.stock_code and h.quantity > 0]
                if codes:
                    websocket_manager.subscribe(cid, codes)
                    from app.services.kis_realtime_stream import kis_stream
                    kis_stream.notify_subscription_change()
                    logger.info(f"[auto-sub] {cid}: 보유종목 {len(codes)}개 자동 구독 → {codes[:5]}...")
        except Exception as e:
            logger.debug(f"[auto-sub] 보유종목 자동 구독 실패: {e}")
    asyncio.create_task(_auto_subscribe_portfolio(client_id))

    try:
        while True:
            # 이 핸들러의 websocket이 active_connections에서 교체되었으면 종료
            if websocket_manager.active_connections.get(client_id) is not websocket:
                logger.debug(f"Client {client_id}: 교체된 연결 감지, 이전 핸들러 종료")
                break

            try:
                data = await asyncio.wait_for(websocket.receive_text(), timeout=60.0)
            except asyncio.TimeoutError:
                # 60초 동안 메시지 없음 → ping 전송으로 연결 유지 확인
                try:
                    await websocket.send_text(_dumps({"type": "ping"}))
                except Exception:
                    break
                continue
            try:
                msg = json.loads(data)
            except (json.JSONDecodeError, AttributeError):
                try:
                    await websocket.send_text(f"Echo: {data}")
                except Exception:
                    break
                continue

            msg_type = msg.get("type", "")

            if msg_type == "ping":
                try:
                    await websocket.send_text(_dumps({"type": "pong"}))
                except Exception:
                    break
                continue

            if msg_type == "subscribe":
                codes = [str(c) for c in msg.get("stock_codes", [])]
                websocket_manager.subscribe(client_id, codes)
                await websocket_manager.send_personal_message(
                    _dumps({
                        "type": "ack", "action": "subscribe",
                        "stock_codes": codes,
                        "subscriptions": websocket_manager.get_subscriptions(client_id),
                    }),
                    client_id,
                )
                continue

            if msg_type == "unsubscribe":
                codes = [str(c) for c in msg.get("stock_codes", [])]
                websocket_manager.unsubscribe(client_id, codes)
                await websocket_manager.send_personal_message(
                    _dumps({
                        "type": "ack", "action": "unsubscribe",
                        "stock_codes": codes,
                        "subscriptions": websocket_manager.get_subscriptions(client_id),
                    }),
                    client_id,
                )
                continue

            if msg_type == "set_subscriptions":
                codes = [str(c) for c in msg.get("stock_codes", [])]
                logger.info(f"[WS] {client_id}: set_subscriptions → {codes}")
                websocket_manager.set_subscriptions(client_id, codes)
                # 구독 목록 변경 → KIS 스트림에 즉시 동기화 요청
                from app.services.kis_realtime_stream import kis_stream
                kis_stream.notify_subscription_change()
                try:
                    await websocket.send_text(_dumps({
                        "type": "ack", "action": "set_subscriptions",
                        "stock_codes": codes,
                        "subscriptions": websocket_manager.get_subscriptions(client_id),
                    }))
                except Exception:
                    break
                # 구독 종목 차트 캐시 사전 워밍업 (백그라운드, 논블로킹)
                if codes:
                    async def _warmup_chart_cache(stock_codes: list):
                        from app.api.endpoints.market_data import get_chart_data_cached, chart_cache
                        from datetime import datetime
                        for code in stock_codes:
                            cached = chart_cache.get(f"{code}__D__200")
                            if cached and (datetime.now() - cached["timestamp"]).total_seconds() < 600:
                                continue
                            try:
                                await get_chart_data_cached(code, "D", 200)
                                logger.debug(f"[warmup] {code} 차트 캐시 완료")
                            except Exception as e:
                                logger.debug(f"[warmup] {code} 캐시 실패: {e}")
                    asyncio.create_task(_warmup_chart_cache(codes))
                continue

            if msg_type == "get_chart":
                stock_code = str(msg.get("stock_code", "")).strip()
                period = msg.get("period", "D")
                limit = int(msg.get("limit", 200))
                if not stock_code:
                    continue
                # 자동 구독 폴백: 클라이언트에 구독 종목이 없으면 요청 종목 자동 구독
                current_subs = websocket_manager.get_subscriptions(client_id)
                if stock_code not in current_subs:
                    websocket_manager.subscribe(client_id, [stock_code])
                    from app.services.kis_realtime_stream import kis_stream
                    kis_stream.notify_subscription_change()
                    logger.info(f"[auto-sub] {client_id}: get_chart 요청으로 {stock_code} 자동 구독 추가")
                try:
                    from app.api.endpoints.market_data import get_chart_data_cached
                    chart_data = await get_chart_data_cached(stock_code, period, limit)
                    logger.info(f"[get_chart] {stock_code}/{period} → {len(chart_data)} candles → sending to {client_id}")
                    await websocket.send_text(
                        _dumps({"type": "chart_data", "stock_code": stock_code, "period": period, "data": chart_data}),
                    )
                except Exception as e:
                    logger.error(f"get_chart 오류 ({stock_code}): {e}")
                    try:
                        await websocket.send_text(
                            _dumps({"type": "chart_data", "stock_code": stock_code, "period": period, "data": [], "error": str(e)}),
                        )
                    except Exception:
                        break
                continue

            await websocket_manager.send_personal_message(f"Echo: {data}", client_id)

    except WebSocketDisconnect as e:
        # 동일 client_id로 새 연결이 교체되었을 경우, 이전 핸들러가 disconnect하면 새 연결까지 제거됨
        # 현재 websocket이 active_connections에 있는 것과 동일한 경우에만 disconnect
        if websocket_manager.active_connections.get(client_id) is websocket:
            websocket_manager.disconnect(client_id)
        logger.info(f"Client {client_id} disconnected (code={e.code})")
    except Exception as e:
        if websocket_manager.active_connections.get(client_id) is websocket:
            websocket_manager.disconnect(client_id)
        logger.error(f"Client {client_id} unexpected WS error: {type(e).__name__}: {e}")

@app.post("/api/admin/reset-daily-loss")
async def reset_daily_loss(api_key: str = Depends(_require_admin_key)):
    """테스트용 일일 손실 한도 강제 리셋 (인증 필수)"""
    from app.services.mdd_manager import mdd_manager
    await mdd_manager.reset()
    return {"status": "ok", "message": "daily loss limit reset"}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app.main:app",
        host=settings.BACKEND_HOST,
        port=settings.BACKEND_PORT,
        reload=settings.DEBUG,
    )