"""
백테스트 엔진

3가지 주요 기능:
1. ML 학습 데이터 생성 (TensorFlow/PyTorch 학습용 라벨링)
2. 전략 성과 백테스트 (과거 수익률 시뮬레이션)
3. 파라미터 최적화 (VectorBT)
"""

import asyncio
import logging
from datetime import date, datetime, timedelta
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, asdict
from decimal import Decimal

import numpy as np
import pandas as pd
from sqlalchemy import select, and_

from app.core.database import AsyncSessionLocal
from app.models.models import StockDailyPrice
from app.services.feature_engineering import (
    db_data_to_dataframe,
    extract_features,
    prepare_training_data,
    FEATURE_COLUMNS,
)

logger = logging.getLogger(__name__)


# ── 데이터 클래스 ────────────────────────────────────────────────────────────

@dataclass
class TrainingDataConfig:
    """ML 학습 데이터 생성 설정"""
    stock_codes: List[str]
    start_date: date
    end_date: date
    forward_days: int = 5  # N일 후 수익률로 라벨링
    up_threshold: float = 0.03  # +3% 이상 = UP (label 2)
    down_threshold: float = -0.03  # -3% 이하 = DOWN (label 0)
    min_samples_per_stock: int = 50  # 종목당 최소 샘플 수


@dataclass
class BacktestConfig:
    """전략 성과 백테스트 설정"""
    stock_codes: List[str]
    start_date: date
    end_date: date
    initial_capital: float = 10_000_000
    strategy_params: Dict[str, Any] = None  # 프리셋 파라미터
    max_positions: int = 5
    position_size: float = 0.2  # 종목당 20%
    commission_rate: float = 0.00015  # 매수/매도 수수료 0.015%
    tax_rate: float = 0.0023  # 매도세 0.23%


@dataclass
class Trade:
    """거래 기록"""
    date: date
    stock_code: str
    action: str  # BUY | SELL
    price: float
    quantity: int
    amount: float
    commission: float
    reason: str  # 진입/청산 이유


@dataclass
class Position:
    """보유 포지션"""
    stock_code: str
    quantity: int
    avg_price: float
    entry_date: date
    entry_reason: str
    current_price: float = 0.0
    unrealized_pnl: float = 0.0
    unrealized_pnl_pct: float = 0.0


@dataclass
class BacktestResult:
    """백테스트 결과"""
    config: Dict[str, Any]
    total_trades: int
    win_trades: int
    lose_trades: int
    win_rate: float
    total_return: float
    total_return_pct: float
    mdd: float  # Maximum Drawdown
    mdd_pct: float
    sharpe_ratio: float
    sortino_ratio: float
    profit_factor: float  # 총이익 / 총손실
    avg_win: float
    avg_lose: float
    max_win: float
    max_lose: float
    equity_curve: List[Dict[str, Any]]  # 일별 자산 추이
    trades: List[Dict[str, Any]]  # 거래 내역
    final_positions: List[Dict[str, Any]]  # 최종 보유 종목


# ── 백테스트 엔진 ─────────────────────────────────────────────────────────────

class BacktestEngine:
    """백테스트 핵심 엔진"""

    def __init__(self):
        self.cash = 0.0
        self.positions: Dict[str, Position] = {}
        self.trades: List[Trade] = []
        self.equity_curve: List[Dict[str, Any]] = []
        self.daily_returns: List[float] = []

    async def _fetch_daily_data(
        self,
        stock_codes: List[str],
        start_date: date,
        end_date: date
    ) -> Dict[str, pd.DataFrame]:
        """DB에서 일봉 데이터 조회"""
        result = {}
        async with AsyncSessionLocal() as session:
            for code in stock_codes:
                stmt = select(StockDailyPrice).where(
                    and_(
                        StockDailyPrice.stock_code == code,
                        StockDailyPrice.trade_date >= start_date,
                        StockDailyPrice.trade_date <= end_date
                    )
                ).order_by(StockDailyPrice.trade_date)
                
                db_result = await session.execute(stmt)
                rows = db_result.scalars().all()
                
                if rows:
                    df = db_data_to_dataframe([{
                        "trade_date": r.trade_date,
                        "open_price": float(r.open_price),
                        "high_price": float(r.high_price),
                        "low_price": float(r.low_price),
                        "close_price": float(r.close_price),
                        "volume": float(r.volume),
                    } for r in rows])
                    result[code] = df
        
        return result

    def _calculate_commission(self, amount: float, is_sell: bool, config: BacktestConfig) -> float:
        """수수료 계산 (매수/매도 수수료 + 매도세)"""
        commission = amount * config.commission_rate
        if is_sell:
            commission += amount * config.tax_rate
        return commission

    def _execute_buy(
        self,
        stock_code: str,
        price: float,
        amount: float,
        current_date: date,
        reason: str,
        config: BacktestConfig
    ):
        """매수 실행"""
        quantity = int(amount / price)
        if quantity == 0:
            return
        
        total_cost = price * quantity
        commission = self._calculate_commission(total_cost, False, config)
        total_amount = total_cost + commission
        
        if total_amount > self.cash:
            return
        
        self.cash -= total_amount
        
        if stock_code in self.positions:
            # 기존 포지션에 추가 매수
            pos = self.positions[stock_code]
            new_quantity = pos.quantity + quantity
            new_avg_price = (pos.avg_price * pos.quantity + price * quantity) / new_quantity
            pos.quantity = new_quantity
            pos.avg_price = new_avg_price
        else:
            self.positions[stock_code] = Position(
                stock_code=stock_code,
                quantity=quantity,
                avg_price=price,
                entry_date=current_date,
                entry_reason=reason
            )
        
        self.trades.append(Trade(
            date=current_date,
            stock_code=stock_code,
            action="BUY",
            price=price,
            quantity=quantity,
            amount=total_cost,
            commission=commission,
            reason=reason
        ))

    def _execute_sell(
        self,
        stock_code: str,
        price: float,
        current_date: date,
        reason: str,
        config: BacktestConfig
    ):
        """매도 실행 (전량)"""
        if stock_code not in self.positions:
            return
        
        pos = self.positions[stock_code]
        total_amount = price * pos.quantity
        commission = self._calculate_commission(total_amount, True, config)
        net_amount = total_amount - commission
        
        self.cash += net_amount
        
        self.trades.append(Trade(
            date=current_date,
            stock_code=stock_code,
            action="SELL",
            price=price,
            quantity=pos.quantity,
            amount=total_amount,
            commission=commission,
            reason=reason
        ))
        
        del self.positions[stock_code]

    def _generate_signals(
        self,
        stock_code: str,
        df: pd.DataFrame,
        features_df: pd.DataFrame,
        current_idx: int,
        config: BacktestConfig
    ) -> Optional[Dict[str, Any]]:
        """전략 신호 생성 (RSI + 볼린저밴드 평균회귀)"""
        if current_idx < 60:  # 충분한 데이터 필요
            return None
        
        # 현재 시점의 피처
        rsi = features_df.loc[current_idx, "rsi_14"]
        bb_pos = features_df.loc[current_idx, "bb_position"]
        ma20_ratio = features_df.loc[current_idx, "ma20_ratio"]
        
        if pd.isna(rsi) or pd.isna(bb_pos) or pd.isna(ma20_ratio):
            return None
        
        # 전략 파라미터 (프리셋 또는 기본값)
        params = config.strategy_params or {}
        rsi_threshold = params.get("MEAN_REVERSION_RSI_THRESHOLD", 30)
        stop_loss = params.get("MEAN_REVERSION_STOP_LOSS_PCT", -6.0) / 100
        take_profit = params.get("MEAN_REVERSION_TAKE_PROFIT_PCT", 12.0) / 100
        rsi_overbought = params.get("RSI_OVERBOUGHT_LEVEL", 70.0)

        current_price = df.loc[current_idx, "close"]

        # 매수 신호 (과매도 + BB 하단)
        if rsi < rsi_threshold and bb_pos < 0.2:
            if stock_code not in self.positions:
                return {
                    "action": "BUY",
                    "price": current_price,
                    "reason": f"RSI={rsi:.1f} BB하단"
                }

        # 매도 신호 (손절/익절)
        if stock_code in self.positions:
            pos = self.positions[stock_code]
            pnl_pct = (current_price - pos.avg_price) / pos.avg_price

            if pnl_pct <= stop_loss:
                return {
                    "action": "SELL",
                    "price": current_price,
                    "reason": f"손절 {pnl_pct*100:.1f}%"
                }
            elif pnl_pct >= take_profit:
                return {
                    "action": "SELL",
                    "price": current_price,
                    "reason": f"익절 {pnl_pct*100:.1f}%"
                }
            elif rsi > rsi_overbought and bb_pos > 0.8:
                return {
                    "action": "SELL",
                    "price": current_price,
                    "reason": f"과매수 청산"
                }
        
        return None

    def _calculate_portfolio_value(self, current_prices: Dict[str, float]) -> float:
        """현재 포트폴리오 가치 계산"""
        position_value = sum(
            pos.quantity * current_prices.get(code, pos.avg_price)
            for code, pos in self.positions.items()
        )
        return self.cash + position_value

    def _calculate_metrics(self, config: BacktestConfig) -> BacktestResult:
        """성과 지표 계산"""
        if not self.trades:
            return BacktestResult(
                config=asdict(config),
                total_trades=0,
                win_trades=0,
                lose_trades=0,
                win_rate=0.0,
                total_return=0.0,
                total_return_pct=0.0,
                mdd=0.0,
                mdd_pct=0.0,
                sharpe_ratio=0.0,
                sortino_ratio=0.0,
                profit_factor=0.0,
                avg_win=0.0,
                avg_lose=0.0,
                max_win=0.0,
                max_lose=0.0,
                equity_curve=self.equity_curve,
                trades=[],
                final_positions=[]
            )
        
        # 거래별 수익률 계산
        buy_trades = {t.stock_code: t for t in self.trades if t.action == "BUY"}
        sell_trades = [t for t in self.trades if t.action == "SELL"]
        
        trade_returns = []
        for sell in sell_trades:
            if sell.stock_code in buy_trades:
                buy = buy_trades[sell.stock_code]
                pnl = (sell.price - buy.price) * sell.quantity - sell.commission - buy.commission
                pnl_pct = pnl / (buy.price * buy.quantity)
                trade_returns.append(pnl_pct)
        
        win_trades = [r for r in trade_returns if r > 0]
        lose_trades = [r for r in trade_returns if r < 0]
        
        # 최종 자산
        final_value = self.equity_curve[-1]["value"] if self.equity_curve else config.initial_capital
        total_return = final_value - config.initial_capital
        total_return_pct = total_return / config.initial_capital
        
        # MDD 계산
        peak = config.initial_capital
        mdd = 0.0
        for point in self.equity_curve:
            if point["value"] > peak:
                peak = point["value"]
            drawdown = point["value"] - peak
            if drawdown < mdd:
                mdd = drawdown
        mdd_pct = mdd / peak if peak > 0 else 0.0
        
        # Sharpe/Sortino
        if len(self.daily_returns) > 1:
            avg_return = np.mean(self.daily_returns)
            std_return = np.std(self.daily_returns, ddof=1)
            downside_returns = [r for r in self.daily_returns if r < 0]
            downside_std = np.std(downside_returns, ddof=1) if downside_returns else 1e-10
            
            sharpe = (avg_return * 252) / (std_return * np.sqrt(252)) if std_return > 0 else 0.0
            sortino = (avg_return * 252) / (downside_std * np.sqrt(252)) if downside_std > 0 else 0.0
        else:
            sharpe = sortino = 0.0
        
        # Profit Factor
        total_win = sum(win_trades) if win_trades else 0.0
        total_lose = abs(sum(lose_trades)) if lose_trades else 1e-10
        profit_factor = total_win / total_lose if total_lose > 0 else 0.0
        
        return BacktestResult(
            config=asdict(config),
            total_trades=len(sell_trades),
            win_trades=len(win_trades),
            lose_trades=len(lose_trades),
            win_rate=len(win_trades) / len(sell_trades) if sell_trades else 0.0,
            total_return=total_return,
            total_return_pct=total_return_pct,
            mdd=mdd,
            mdd_pct=mdd_pct,
            sharpe_ratio=sharpe,
            sortino_ratio=sortino,
            profit_factor=profit_factor,
            avg_win=np.mean(win_trades) if win_trades else 0.0,
            avg_lose=np.mean(lose_trades) if lose_trades else 0.0,
            max_win=max(win_trades) if win_trades else 0.0,
            max_lose=min(lose_trades) if lose_trades else 0.0,
            equity_curve=self.equity_curve,
            trades=[asdict(t) for t in self.trades],
            final_positions=[asdict(p) for p in self.positions.values()]
        )

    async def run_backtest(self, config: BacktestConfig) -> BacktestResult:
        """전략 성과 백테스트 실행"""
        logger.info(f"백테스트 시작: {config.start_date} ~ {config.end_date}, {len(config.stock_codes)}종목")
        
        self.cash = config.initial_capital
        self.positions = {}
        self.trades = []
        self.equity_curve = []
        self.daily_returns = []
        
        # 데이터 로드
        stock_data = await self._fetch_daily_data(
            config.stock_codes,
            config.start_date,
            config.end_date
        )
        
        if not stock_data:
            logger.warning("데이터 없음")
            return self._calculate_metrics(config)
        
        # 피처 추출
        features_data = {}
        for code, df in stock_data.items():
            features_df = extract_features(df)
            features_data[code] = features_df
        
        # 날짜별 시뮬레이션
        all_dates = sorted(set(
            date_val for df in stock_data.values()
            for date_val in df["date"].values
        ))
        
        prev_portfolio_value = config.initial_capital
        
        for current_date in all_dates:
            fmt = "%Y%m%d" if len(current_date) == 8 else "%Y-%m-%d"
            current_date_obj = datetime.strptime(current_date, fmt).date()
            
            # 각 종목별 신호 생성 및 주문 실행
            current_prices = {}
            for code, df in stock_data.items():
                date_mask = df["date"] == current_date
                if not date_mask.any():
                    continue
                
                current_idx = df.index[date_mask][0]
                current_price = df.loc[current_idx, "close"]
                current_prices[code] = current_price
                
                # 신호 생성
                signal = self._generate_signals(
                    code, df, features_data[code], current_idx, config
                )
                
                if signal:
                    if signal["action"] == "BUY":
                        # 포지션 크기 제한
                        if len(self.positions) < config.max_positions:
                            position_amount = self.cash * config.position_size
                            self._execute_buy(
                                code, signal["price"], position_amount,
                                current_date_obj, signal["reason"], config
                            )
                    elif signal["action"] == "SELL":
                        self._execute_sell(
                            code, signal["price"], current_date_obj,
                            signal["reason"], config
                        )
            
            # 일별 자산 기록
            portfolio_value = self._calculate_portfolio_value(current_prices)
            daily_return = (portfolio_value - prev_portfolio_value) / prev_portfolio_value
            self.daily_returns.append(daily_return)
            prev_portfolio_value = portfolio_value
            
            self.equity_curve.append({
                "date": current_date,
                "value": portfolio_value,
                "cash": self.cash,
                "positions": len(self.positions),
                "return_pct": (portfolio_value - config.initial_capital) / config.initial_capital
            })
        
        result = self._calculate_metrics(config)
        logger.info(f"백테스트 완료: 수익률 {result.total_return_pct*100:.2f}%, 승률 {result.win_rate*100:.1f}%")
        return result


# ── ML 학습 데이터 생성 ────────────────────────────────────────────────────────

async def generate_training_data(config: TrainingDataConfig) -> Dict[str, Any]:
    """ML 학습 데이터 생성 (라벨링된 피처)
    
    Returns:
        {
            "total_samples": int,
            "label_distribution": {"UP": int, "HOLD": int, "DOWN": int},
            "data": pd.DataFrame with columns [date, stock_code, features..., label]
        }
    """
    logger.info(f"ML 학습 데이터 생성 시작: {len(config.stock_codes)}종목")
    
    engine = BacktestEngine()
    stock_data = await engine._fetch_daily_data(
        config.stock_codes,
        config.start_date,
        config.end_date
    )
    
    all_samples = []
    
    for stock_code, df in stock_data.items():
        if len(df) < config.forward_days + 60:
            continue
        
        # 피처 추출
        features_df = extract_features(df)
        
        # 라벨 생성
        X_tabular, X_seq, y, valid_indices = prepare_training_data(
            df, features_df,
            forecast_days=config.forward_days,
            up_threshold=config.up_threshold,
            down_threshold=config.down_threshold
        )
        
        if len(valid_indices) < config.min_samples_per_stock:
            logger.warning(f"{stock_code}: 샘플 부족 ({len(valid_indices)}개)")
            continue
        
        # DataFrame으로 변환
        for i, idx in enumerate(valid_indices):
            sample = {
                "date": df.loc[idx, "date"],
                "stock_code": stock_code,
                **{col: X_tabular[i, j] for j, col in enumerate(FEATURE_COLUMNS)},
                "label": int(y[i])
            }
            all_samples.append(sample)
    
    result_df = pd.DataFrame(all_samples)
    
    if len(result_df) == 0:
        return {
            "error": "학습 데이터 생성 실패",
            "total_samples": 0,
            "label_distribution": {}
        }
    
    label_dist = result_df["label"].value_counts().to_dict()
    label_names = {0: "DOWN", 1: "HOLD", 2: "UP"}
    
    logger.info(f"ML 학습 데이터 생성 완료: {len(result_df)}샘플, 분포 {label_dist}")
    
    return {
        "total_samples": len(result_df),
        "label_distribution": {label_names[k]: int(v) for k, v in label_dist.items()},
        "data": result_df,
        "config": asdict(config)
    }