"""Work/learning trade-off balance for trading agents.

This module provides the WorkTradeBalance class for deciding whether
trading agents should work (trade) or learn based on economic status,
skill level, and win rate.
"""

from dataclasses import dataclass
from enum import Enum
from typing import Optional

from openclaw.core.economy import SurvivalStatus, TradingEconomicTracker
from openclaw.utils.logging import get_logger


class ActivityDecision(str, Enum):
    """Decision outcome for agent activity."""

    TRADE = "trade"
    LEARN = "learn"
    MINIMAL_TRADE = "minimal_trade"
    PAPER_TRADE = "paper_trade"


class LearningPriority(str, Enum):
    """Priority level for learning activities."""

    ALL_COURSES = "all_courses"
    ESSENTIAL_ONLY = "essential_only"
    EMERGENCY_ONLY = "emergency_only"
    NONE = "none"


@dataclass
class TradeIntensity:
    """Trade intensity configuration based on win rate."""

    position_size_multiplier: float
    max_concurrent_positions: int
    risk_per_trade: float


@dataclass
class WorkTradeConfig:
    """Configuration for work/learning balance decisions."""

    # Base activity weights by economic status (must sum to 1.0)
    thriving_weights: dict[str, float] = None
    stable_weights: dict[str, float] = None
    struggling_weights: dict[str, float] = None
    critical_weights: dict[str, float] = None

    # Skill adjustment thresholds
    low_skill_threshold: float = 0.3
    high_skill_threshold: float = 0.8
    low_skill_learning_boost: float = 0.15

    # Win rate thresholds for trade intensity
    high_win_rate_threshold: float = 0.6
    low_win_rate_threshold: float = 0.4

    # Trade intensity by win rate
    aggressive_intensity: TradeIntensity = None
    normal_intensity: TradeIntensity = None
    conservative_intensity: TradeIntensity = None
    minimal_intensity: TradeIntensity = None

    def __post_init__(self):
        """Set default values for None fields."""
        if self.thriving_weights is None:
            self.thriving_weights = {"trade": 0.70, "learn": 0.30}
        if self.stable_weights is None:
            self.stable_weights = {"trade": 0.80, "learn": 0.20}
        if self.struggling_weights is None:
            self.struggling_weights = {"trade": 0.90, "learn": 0.10}
        if self.critical_weights is None:
            self.critical_weights = {"trade": 1.00, "learn": 0.00}

        if self.aggressive_intensity is None:
            self.aggressive_intensity = TradeIntensity(
                position_size_multiplier=1.5,
                max_concurrent_positions=5,
                risk_per_trade=0.03,
            )
        if self.normal_intensity is None:
            self.normal_intensity = TradeIntensity(
                position_size_multiplier=1.0,
                max_concurrent_positions=3,
                risk_per_trade=0.02,
            )
        if self.conservative_intensity is None:
            self.conservative_intensity = TradeIntensity(
                position_size_multiplier=0.6,
                max_concurrent_positions=2,
                risk_per_trade=0.01,
            )
        if self.minimal_intensity is None:
            self.minimal_intensity = TradeIntensity(
                position_size_multiplier=0.3,
                max_concurrent_positions=1,
                risk_per_trade=0.005,
            )


class WorkTradeBalance:
    """Balance work (trading) and learning based on economic status.

    This class helps agents decide whether to trade or learn based on:
    - Economic status (thriving, stable, struggling, critical)
    - Skill level (lower skill = more learning)
    - Win rate (lower win rate = more learning/conservative trading)

    Args:
        economic_tracker: The agent's economic tracker
        config: Configuration for decision logic (optional)
    """

    def __init__(
        self,
        economic_tracker: TradingEconomicTracker,
        config: Optional[WorkTradeConfig] = None,
    ):
        self.economic_tracker = economic_tracker
        self.config = config or WorkTradeConfig()
        self.logger = get_logger(f"work_trade_balance.{economic_tracker.agent_id}")

    def decide_activity(
        self,
        skill_level: float,
        win_rate: float,
        random_draw: float = 0.5,
    ) -> ActivityDecision:
        """Decide whether to trade or learn.

        Args:
            skill_level: Current skill level (0.0 to 1.0)
            win_rate: Current win rate (0.0 to 1.0)
            random_draw: Random value 0-1 for probabilistic decision

        Returns:
            ActivityDecision enum value
        """
        status = self.economic_tracker.get_survival_status()

        # Critical status: only minimal trading, no learning
        if status == SurvivalStatus.CRITICAL:
            self.logger.warning(
                f"Status {status.value}: FORCED minimal trading only"
            )
            return ActivityDecision.MINIMAL_TRADE

        # Bankrupt: should not reach here, but handle gracefully
        if status == SurvivalStatus.BANKRUPT:
            self.logger.error("Status bankrupt: cannot perform any activity")
            return ActivityDecision.MINIMAL_TRADE

        # Get base weights for current status
        weights = self._get_weights_for_status(status)

        # Adjust weights based on skill level
        adjusted_weights = self._adjust_for_skill(weights, skill_level, status)

        # Decide based on adjusted weights
        decision = self._make_decision(adjusted_weights, random_draw)

        # Log decision
        self.logger.info(
            f"Status={status.value}, skill={skill_level:.1%}, win_rate={win_rate:.1%}, "
            f"decision={decision.value}"
        )

        return decision

    def get_learning_priority(self) -> LearningPriority:
        """Determine what type of learning is appropriate.

        Returns:
            LearningPriority enum value
        """
        status = self.economic_tracker.get_survival_status()

        if status == SurvivalStatus.THRIVING:
            return LearningPriority.ALL_COURSES
        elif status == SurvivalStatus.STABLE:
            return LearningPriority.ESSENTIAL_ONLY
        elif status == SurvivalStatus.STRUGGLING:
            return LearningPriority.EMERGENCY_ONLY
        else:
            return LearningPriority.NONE

    def get_trade_intensity(self, win_rate: float) -> TradeIntensity:
        """Get trade intensity configuration based on win rate.

        Args:
            win_rate: Current win rate (0.0 to 1.0)

        Returns:
            TradeIntensity configuration
        """
        status = self.economic_tracker.get_survival_status()

        # Critical status: always minimal
        if status in (SurvivalStatus.CRITICAL, SurvivalStatus.BANKRUPT):
            return self.config.minimal_intensity

        # Struggling: conservative regardless of win rate
        if status == SurvivalStatus.STRUGGLING:
            return self.config.conservative_intensity

        # Based on win rate
        if win_rate >= self.config.high_win_rate_threshold:
            return self.config.aggressive_intensity
        elif win_rate >= self.config.low_win_rate_threshold:
            return self.config.normal_intensity
        else:
            return self.config.conservative_intensity

    def should_paper_trade(
        self,
        skill_level: float,
        win_rate: float,
        consecutive_losses: int = 0,
    ) -> bool:
        """Decide if agent should use paper trading instead of real trading.

        Paper trading is recommended when:
        - Skill level is very low
        - Win rate is very low
        - Agent has consecutive losses

        Args:
            skill_level: Current skill level
            win_rate: Current win rate
            consecutive_losses: Number of consecutive losing trades

        Returns:
            True if paper trading is recommended
        """
        # Very low skill or win rate
        if skill_level < 0.2 or win_rate < 0.3:
            return True

        # Consecutive losses threshold
        if consecutive_losses >= 3:
            return True

        # Struggling status with poor performance
        status = self.economic_tracker.get_survival_status()
        if status == SurvivalStatus.STRUGGLING and win_rate < 0.45:
            return True

        return False

    def get_learning_investment_budget(self) -> float:
        """Calculate how much can be spent on learning.

        Returns:
            Maximum amount to invest in learning as % of balance
        """
        status = self.economic_tracker.get_survival_status()
        balance = self.economic_tracker.balance

        if status == SurvivalStatus.THRIVING:
            return balance * 0.10  # 10% of balance
        elif status == SurvivalStatus.STABLE:
            return balance * 0.05  # 5% of balance
        elif status == SurvivalStatus.STRUGGLING:
            return balance * 0.02  # 2% of balance
        else:
            return 0.0  # No learning budget

    def _get_weights_for_status(self, status: SurvivalStatus) -> dict[str, float]:
        """Get base activity weights for survival status."""
        if status == SurvivalStatus.THRIVING:
            return self.config.thriving_weights.copy()
        elif status == SurvivalStatus.STABLE:
            return self.config.stable_weights.copy()
        elif status == SurvivalStatus.STRUGGLING:
            return self.config.struggling_weights.copy()
        else:
            return self.config.critical_weights.copy()

    def _adjust_for_skill(
        self,
        weights: dict[str, float],
        skill_level: float,
        status: SurvivalStatus,
    ) -> dict[str, float]:
        """Adjust weights based on skill level.

        Lower skill = more learning weight (unless critical/struggling)
        """
        adjusted = weights.copy()

        # Don't adjust if in critical or struggling status
        if status in (SurvivalStatus.CRITICAL, SurvivalStatus.STRUGGLING):
            return adjusted

        # Low skill: boost learning weight
        if skill_level < self.config.low_skill_threshold:
            boost = self.config.low_skill_learning_boost
            adjusted["learn"] = min(0.5, adjusted["learn"] + boost)
            adjusted["trade"] = 1.0 - adjusted["learn"]

        # High skill: slightly reduce learning (optional optimization)
        elif skill_level > self.config.high_skill_threshold:
            # Slight preference for trading when highly skilled
            reduction = 0.05
            adjusted["learn"] = max(0.05, adjusted["learn"] - reduction)
            adjusted["trade"] = 1.0 - adjusted["learn"]

        return adjusted

    def _make_decision(
        self,
        weights: dict[str, float],
        random_draw: float,
    ) -> ActivityDecision:
        """Make activity decision based on weights and random draw."""
        if random_draw < weights["trade"]:
            return ActivityDecision.TRADE
        else:
            return ActivityDecision.LEARN

    def get_decision_summary(
        self,
        skill_level: float,
        win_rate: float,
    ) -> dict:
        """Get comprehensive decision summary for logging/debugging.

        Args:
            skill_level: Current skill level
            win_rate: Current win rate

        Returns:
            Dictionary with all decision factors
        """
        status = self.economic_tracker.get_survival_status()
        weights = self._get_weights_for_status(status)
        adjusted = self._adjust_for_skill(weights, skill_level, status)
        intensity = self.get_trade_intensity(win_rate)

        return {
            "status": status.value,
            "balance": self.economic_tracker.balance,
            "skill_level": skill_level,
            "win_rate": win_rate,
            "base_weights": weights,
            "adjusted_weights": adjusted,
            "learning_priority": self.get_learning_priority().value,
            "trade_intensity": {
                "position_size_multiplier": intensity.position_size_multiplier,
                "max_concurrent_positions": intensity.max_concurrent_positions,
                "risk_per_trade": intensity.risk_per_trade,
            },
            "learning_budget": self.get_learning_investment_budget(),
            "should_paper_trade": self.should_paper_trade(skill_level, win_rate),
        }
