"""
风险预算系统

月度风险预算分配与追踪
"""

from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Dict, List, Optional
from enum import Enum


class BudgetStatus(Enum):
    """预算状态"""
    NORMAL = "normal"
    WARNING = "warning"
    COOLING = "cooling"
    STOPPED = "stopped"


@dataclass
class BudgetState:
    """预算状态"""
    monthly_budget: float
    consumed: float
    remaining: float
    status: BudgetStatus
    cooling_end_date: Optional[datetime] = None


class RiskBudgetManager:
    """
    风险预算管理器

    功能：
    1. 月度预算分配（默认账户价值的5%）
    2. 单笔交易风险消耗追踪
    3. 触发阈值管理（预警10%、冷却20%、停止50%）
    """

    def __init__(
        self,
        monthly_budget_pct: float = 0.05,
        warning_threshold: float = 0.10,
        cooling_threshold: float = 0.20,
        stop_threshold: float = 0.50,
        cooling_period_days: int = 3
    ):
        self.monthly_budget_pct = monthly_budget_pct
        self.warning_threshold = warning_threshold
        self.cooling_threshold = cooling_threshold
        self.stop_threshold = stop_threshold
        self.cooling_period_days = cooling_period_days

        self.current_budget: Optional[BudgetState] = None
        self.trade_history: List[Dict] = []

    def allocate_monthly_budget(self, account_value: float):
        """分配月度预算"""
        budget = account_value * self.monthly_budget_pct

        self.current_budget = BudgetState(
            monthly_budget=budget,
            consumed=0.0,
            remaining=budget,
            status=BudgetStatus.NORMAL
        )

    def record_trade_result(
        self,
        trade_id: str,
        realized_loss: float,
        timestamp: Optional[datetime] = None
    ) -> BudgetState:
        """
        记录交易结果，更新预算消耗

        Args:
            trade_id: 交易ID
            realized_loss: 实际亏损（正数）
            timestamp: 时间戳

        Returns:
            BudgetState: 更新后的预算状态
        """
        if self.current_budget is None:
            raise RuntimeError("Budget not allocated. Call allocate_monthly_budget first.")

        if timestamp is None:
            timestamp = datetime.now()

        # 只记录亏损
        if realized_loss > 0:
            self.current_budget.consumed += realized_loss
            self.current_budget.remaining = max(0, self.current_budget.monthly_budget - self.current_budget.consumed)

            self.trade_history.append({
                "trade_id": trade_id,
                "loss": realized_loss,
                "timestamp": timestamp,
                "cumulative_consumed": self.current_budget.consumed
            })

        # 更新状态
        consumption_ratio = self.current_budget.consumed / self.current_budget.monthly_budget

        if consumption_ratio >= self.stop_threshold:
            self.current_budget.status = BudgetStatus.STOPPED
        elif consumption_ratio >= self.cooling_threshold:
            if self.current_budget.status != BudgetStatus.COOLING:
                self.current_budget.status = BudgetStatus.COOLING
                self.current_budget.cooling_end_date = timestamp + timedelta(days=self.cooling_period_days)
        elif consumption_ratio >= self.warning_threshold:
            self.current_budget.status = BudgetStatus.WARNING

        return self.current_budget

    def check_can_trade(self) -> bool:
        """检查是否可以交易 - 激进模式：除非完全停止，否则允许交易"""
        if self.current_budget is None:
            return True

        # 检查冷却期
        if self.current_budget.status == BudgetStatus.COOLING:
            if self.current_budget.cooling_end_date and datetime.now() < self.current_budget.cooling_end_date:
                # 冷却期允许50%仓位
                return True
            else:
                # 冷却期结束，恢复
                self.current_budget.status = BudgetStatus.WARNING
                return True

        # 即使是STOPPED状态，也允许25%仓位交易
        return True

    def get_budget_summary(self) -> Dict:
        """获取预算摘要"""
        if self.current_budget is None:
            return {"status": "not_allocated"}

        return {
            "monthly_budget": self.current_budget.monthly_budget,
            "consumed": self.current_budget.consumed,
            "remaining": self.current_budget.remaining,
            "consumption_ratio": self.current_budget.consumed / self.current_budget.monthly_budget,
            "status": self.current_budget.status.value,
            "can_trade": self.check_can_trade()
        }