"""Strategy portfolio management for combining multiple trading strategies.

This module provides the StrategyPortfolio class for managing multiple strategies
with weight allocation, signal aggregation, and rebalancing capabilities.
"""

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum, auto
from typing import Any, Callable, Dict, List, Optional, Protocol, Type

import numpy as np
import pandas as pd
from pydantic import BaseModel, Field

from openclaw.portfolio.weights import (
    WeightMethod,
    calculate_equal_weights,
    calculate_inverse_volatility_weights,
    calculate_momentum_weights,
    calculate_risk_parity_weights,
    normalize_weights,
)
from openclaw.portfolio.signal_aggregator import (
    AggregationMethod,
    AggregatedSignal,
    SignalAggregator,
    StrategySignal,
)
from openclaw.portfolio.rebalancer import (
    RebalanceResult,
    RebalanceTrigger,
    Rebalancer,
    TransactionCostModel,
)
from openclaw.utils.logging import get_logger


class StrategyStatus(str, Enum):
    """Status of a strategy in the portfolio."""

    ACTIVE = "active"
    INACTIVE = "inactive"
    PAUSED = "paused"
    DISABLED = "disabled"


@dataclass
class StrategyConfig:
    """Configuration for a strategy in the portfolio.

    Attributes:
        strategy_id: Unique identifier for the strategy
        strategy_class: Strategy class (optional)
        weight: Target weight in portfolio
        max_weight: Maximum allowed weight
        min_weight: Minimum allowed weight
        status: Current strategy status
        params: Strategy-specific parameters
        metadata: Additional metadata
    """

    strategy_id: str
    strategy_class: Optional[Type[Any]] = None
    weight: float = 0.0
    max_weight: float = 1.0
    min_weight: float = 0.0
    status: StrategyStatus = StrategyStatus.ACTIVE
    params: Dict[str, Any] = field(default_factory=dict)
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class StrategyPerformance:
    """Performance metrics for a strategy.

    Attributes:
        strategy_id: Strategy identifier
        total_return: Total return percentage
        sharpe_ratio: Sharpe ratio
        max_drawdown: Maximum drawdown percentage
        win_rate: Win rate percentage
        volatility: Strategy volatility
        last_updated: Last update timestamp
    """

    strategy_id: str
    total_return: float = 0.0
    sharpe_ratio: float = 0.0
    max_drawdown: float = 0.0
    win_rate: float = 0.0
    volatility: float = 0.0
    trade_count: int = 0
    last_updated: Optional[datetime] = None


@dataclass
class PortfolioState:
    """Current state of the strategy portfolio.

    Attributes:
        portfolio_id: Portfolio identifier
        strategies: Active strategy configurations
        current_weights: Current portfolio weights
        target_weights: Target portfolio weights
        last_rebalance: Last rebalance timestamp
        total_value: Total portfolio value
        cash: Available cash
    """

    portfolio_id: str
    strategies: Dict[str, StrategyConfig] = field(default_factory=dict)
    current_weights: Dict[str, float] = field(default_factory=dict)
    target_weights: Dict[str, float] = field(default_factory=dict)
    last_rebalance: Optional[datetime] = None
    total_value: float = 0.0
    cash: float = 0.0


class StrategyProtocol(Protocol):
    """Protocol for strategies that can be added to portfolio."""

    def generate_signal(self, data: pd.DataFrame) -> Dict[str, Any]:
        """Generate trading signal.

        Args:
            data: Market data

        Returns:
            Signal dictionary with 'signal' and optional 'confidence'
        """
        ...

    def get_performance(self) -> Dict[str, float]:
        """Get strategy performance metrics."""
        ...


class StrategyPortfolio:
    """Multi-strategy portfolio with weight management and signal aggregation.

    This class manages a collection of trading strategies with dynamic weight
    allocation, signal aggregation, and automatic rebalancing.

    Args:
        portfolio_id: Unique identifier for this portfolio
        weight_method: Default weight allocation method
        aggregation_method: Default signal aggregation method
        rebalance_trigger: Rebalancing trigger type
        initial_capital: Initial capital allocation
    """

    def __init__(
        self,
        portfolio_id: str,
        weight_method: WeightMethod = WeightMethod.EQUAL,
        aggregation_method: AggregationMethod = AggregationMethod.WEIGHTED,
        rebalance_trigger: RebalanceTrigger = RebalanceTrigger.PERIODIC,
        initial_capital: float = 100000.0,
    ):
        self.portfolio_id = portfolio_id
        self.weight_method = weight_method
        self.aggregation_method = aggregation_method
        self.initial_capital = initial_capital

        # Strategy management
        self._strategies: Dict[str, StrategyConfig] = {}
        self._strategy_instances: Dict[str, Any] = {}
        self._performance: Dict[str, StrategyPerformance] = {}

        # Weight management
        self._current_weights: Dict[str, float] = {}
        self._target_weights: Dict[str, float] = {}

        # Components
        self._signal_aggregator = SignalAggregator(
            method=aggregation_method,
            confidence_threshold=0.5,
        )
        self._rebalancer = Rebalancer(
            trigger_type=rebalance_trigger,
            rebalance_frequency=30,
            drift_threshold=0.05,
        )

        # State
        self._state = PortfolioState(
            portfolio_id=portfolio_id,
            total_value=initial_capital,
            cash=initial_capital,
        )

        # Returns history for weight calculations
        self._returns_history: pd.DataFrame = pd.DataFrame()

        self.logger = get_logger(f"portfolio.{portfolio_id}")
        self.logger.info(
            f"StrategyPortfolio '{portfolio_id}' initialized with "
            f"${initial_capital:,.2f} capital"
        )

    # ==================== Strategy Lifecycle Management ====================

    def add_strategy(
        self,
        strategy_id: str,
        strategy: Optional[Any] = None,
        weight: float = 0.0,
        params: Optional[Dict[str, Any]] = None,
        status: StrategyStatus = StrategyStatus.ACTIVE,
    ) -> bool:
        """Add a strategy to the portfolio.

        Args:
            strategy_id: Unique identifier for the strategy
            strategy: Strategy instance or class
            weight: Initial weight (0 = auto-calculate)
            params: Strategy parameters
            status: Initial strategy status

        Returns:
            True if strategy was added successfully
        """
        if strategy_id in self._strategies:
            self.logger.warning(f"Strategy '{strategy_id}' already exists")
            return False

        config = StrategyConfig(
            strategy_id=strategy_id,
            weight=weight,
            status=status,
            params=params or {},
        )

        self._strategies[strategy_id] = config
        if strategy is not None:
            self._strategy_instances[strategy_id] = strategy

        # Initialize performance tracking
        self._performance[strategy_id] = StrategyPerformance(
            strategy_id=strategy_id,
        )

        # Recalculate weights if using auto-allocation
        if weight == 0.0:
            self._recalculate_weights()

        self.logger.info(f"Added strategy '{strategy_id}' with status {status.value}")
        return True

    def remove_strategy(self, strategy_id: str) -> bool:
        """Remove a strategy from the portfolio.

        Args:
            strategy_id: Strategy identifier

        Returns:
            True if strategy was removed
        """
        if strategy_id not in self._strategies:
            self.logger.warning(f"Strategy '{strategy_id}' not found")
            return False

        del self._strategies[strategy_id]
        if strategy_id in self._strategy_instances:
            del self._strategy_instances[strategy_id]
        if strategy_id in self._performance:
            del self._performance[strategy_id]
        if strategy_id in self._current_weights:
            del self._current_weights[strategy_id]
        if strategy_id in self._target_weights:
            del self._target_weights[strategy_id]

        # Recalculate weights
        self._recalculate_weights()

        self.logger.info(f"Removed strategy '{strategy_id}'")
        return True

    def enable_strategy(self, strategy_id: str) -> bool:
        """Enable an inactive or paused strategy.

        Args:
            strategy_id: Strategy identifier

        Returns:
            True if strategy was enabled
        """
        if strategy_id not in self._strategies:
            return False

        self._strategies[strategy_id].status = StrategyStatus.ACTIVE
        self.logger.info(f"Enabled strategy '{strategy_id}'")
        return True

    def disable_strategy(self, strategy_id: str) -> bool:
        """Disable a strategy (excluded from signal aggregation).

        Args:
            strategy_id: Strategy identifier

        Returns:
            True if strategy was disabled
        """
        if strategy_id not in self._strategies:
            return False

        self._strategies[strategy_id].status = StrategyStatus.DISABLED
        self._current_weights[strategy_id] = 0.0
        self.logger.info(f"Disabled strategy '{strategy_id}'")
        return True

    def pause_strategy(self, strategy_id: str) -> bool:
        """Pause a strategy temporarily.

        Args:
            strategy_id: Strategy identifier

        Returns:
            True if strategy was paused
        """
        if strategy_id not in self._strategies:
            return False

        self._strategies[strategy_id].status = StrategyStatus.PAUSED
        self.logger.info(f"Paused strategy '{strategy_id}'")
        return True

    def get_strategy_status(self, strategy_id: str) -> Optional[StrategyStatus]:
        """Get the status of a strategy.

        Args:
            strategy_id: Strategy identifier

        Returns:
            Strategy status or None if not found
        """
        config = self._strategies.get(strategy_id)
        return config.status if config else None

    def list_strategies(self, active_only: bool = False) -> List[str]:
        """List all strategies in the portfolio.

        Args:
            active_only: Only return active strategies

        Returns:
            List of strategy identifiers
        """
        if active_only:
            return [
                sid for sid, config in self._strategies.items()
                if config.status == StrategyStatus.ACTIVE
            ]
        return list(self._strategies.keys())

    # ==================== Weight Management ====================

    def _recalculate_weights(self) -> None:
        """Recalculate portfolio weights based on current method."""
        active_strategies = self.list_strategies(active_only=True)

        if not active_strategies:
            self._target_weights = {}
            return

        # Get returns data for weight calculations
        returns_data = self._returns_history if not self._returns_history.empty else None

        # Calculate weights based on method
        if self.weight_method == WeightMethod.EQUAL:
            weights = calculate_equal_weights(active_strategies)
        elif self.weight_method == WeightMethod.RISK_PARITY:
            weights = calculate_risk_parity_weights(active_strategies, returns_data)
        elif self.weight_method == WeightMethod.MOMENTUM:
            weights = calculate_momentum_weights(active_strategies, returns_data)
        elif self.weight_method == WeightMethod.INVERSE_VOLATILITY:
            weights = calculate_inverse_volatility_weights(active_strategies, returns_data)
        else:
            weights = calculate_equal_weights(active_strategies)

        # Apply strategy constraints
        for strategy_id in active_strategies:
            config = self._strategies[strategy_id]
            if strategy_id in weights:
                weights[strategy_id] = max(
                    config.min_weight,
                    min(config.max_weight, weights[strategy_id])
                )

        # Normalize to ensure sum = 1.0
        self._target_weights = normalize_weights(weights)

        # Initialize current weights if empty
        if not self._current_weights:
            self._current_weights = self._target_weights.copy()

    def set_weights(self, weights: Dict[str, float]) -> None:
        """Set custom weights for strategies.

        Args:
            weights: Dictionary mapping strategy_id to weight
        """
        self._target_weights = normalize_weights(weights)
        self.weight_method = WeightMethod.CUSTOM
        self.logger.info(f"Set custom weights: {self._target_weights}")

    def get_weights(self) -> Dict[str, float]:
        """Get current portfolio weights.

        Returns:
            Dictionary of strategy weights
        """
        return self._current_weights.copy()

    def get_target_weights(self) -> Dict[str, float]:
        """Get target portfolio weights.

        Returns:
            Dictionary of target strategy weights
        """
        return self._target_weights.copy()

    def update_weight_method(self, method: WeightMethod) -> None:
        """Update the weight allocation method.

        Args:
            method: New weight method
        """
        self.weight_method = method
        self._recalculate_weights()
        self.logger.info(f"Updated weight method to {method.value}")

    # ==================== Signal Aggregation ====================

    def aggregate_signals(
        self,
        signals: Dict[str, Dict[str, Any]],
        method: Optional[AggregationMethod] = None,
    ) -> AggregatedSignal:
        """Aggregate signals from multiple strategies.

        Args:
            signals: Dictionary mapping strategy_id to signal data
            method: Override default aggregation method

        Returns:
            Aggregated signal result
        """
        strategy_signals = []

        for strategy_id, signal_data in signals.items():
            # Skip disabled strategies
            if strategy_id in self._strategies:
                if self._strategies[strategy_id].status != StrategyStatus.ACTIVE:
                    continue

            signal_str = signal_data.get("signal", "hold")
            confidence = signal_data.get("confidence", 0.5)

            strategy_signals.append(
                StrategySignal(
                    strategy_id=strategy_id,
                    signal=signal_str,
                    confidence=confidence,
                    metadata={k: v for k, v in signal_data.items() if k not in ("signal", "confidence")},
                )
            )

        # Get current weights for weighted aggregation
        weights = self._current_weights

        return self._signal_aggregator.aggregate(
            strategy_signals,
            method=method,
            weights=weights,
        )

    def set_aggregation_method(self, method: AggregationMethod) -> None:
        """Set the default signal aggregation method.

        Args:
            method: New aggregation method
        """
        self.aggregation_method = method
        self._signal_aggregator.method = method
        self.logger.info(f"Updated aggregation method to {method.value}")

    def set_confidence_threshold(self, threshold: float) -> None:
        """Set the confidence threshold for signal aggregation.

        Args:
            threshold: Confidence threshold (0.0 to 1.0)
        """
        self._signal_aggregator.set_confidence_threshold(threshold)

    # ==================== Rebalancing ====================

    def check_rebalance(self, current_time: Optional[datetime] = None) -> bool:
        """Check if portfolio rebalancing is needed.

        Args:
            current_time: Current timestamp

        Returns:
            True if rebalancing is needed
        """
        needed, _ = self._rebalancer.check_rebalance_needed(
            self._current_weights,
            self._target_weights,
            current_time,
        )
        return needed

    def rebalance(
        self,
        current_time: Optional[datetime] = None,
        force: bool = False,
    ) -> Optional[RebalanceResult]:
        """Execute portfolio rebalancing.

        Args:
            current_time: Current timestamp
            force: Force rebalance even if not triggered

        Returns:
            RebalanceResult if rebalancing occurred
        """
        # Ensure we have up-to-date target weights
        if self.weight_method != WeightMethod.CUSTOM:
            self._recalculate_weights()

        result = self._rebalancer.rebalance(
            current_weights=self._current_weights,
            target_weights=self._target_weights,
            portfolio_value=self._state.total_value,
            current_time=current_time,
            force=force,
        )

        if result:
            self._current_weights = result.new_weights.copy()
            self._state.last_rebalance = result.timestamp
            self.logger.info(
                f"Rebalanced portfolio: {result.trades_executed} trades, "
                f"${result.transaction_costs:.2f} costs"
            )

        return result

    def set_rebalance_config(
        self,
        trigger: Optional[RebalanceTrigger] = None,
        frequency: Optional[int] = None,
        drift_threshold: Optional[float] = None,
    ) -> None:
        """Configure rebalancing parameters.

        Args:
            trigger: Rebalance trigger type
            frequency: Days between periodic rebalances
            drift_threshold: Drift threshold for threshold-based rebalancing
        """
        self._rebalancer.update_config(
            trigger_type=trigger,
            rebalance_frequency=frequency,
            drift_threshold=drift_threshold,
        )
        if trigger:
            self.logger.info(f"Updated rebalance trigger to {trigger.value}")

    def get_rebalance_history(self) -> List[RebalanceResult]:
        """Get rebalancing history.

        Returns:
            List of rebalance results
        """
        return self._rebalancer.get_rebalance_history()

    # ==================== Performance Tracking ====================

    def update_performance(
        self,
        strategy_id: str,
        metrics: Dict[str, float],
    ) -> bool:
        """Update performance metrics for a strategy.

        Args:
            strategy_id: Strategy identifier
            metrics: Performance metrics dictionary

        Returns:
            True if updated successfully
        """
        if strategy_id not in self._performance:
            return False

        perf = self._performance[strategy_id]
        perf.total_return = metrics.get("total_return", perf.total_return)
        perf.sharpe_ratio = metrics.get("sharpe_ratio", perf.sharpe_ratio)
        perf.max_drawdown = metrics.get("max_drawdown", perf.max_drawdown)
        perf.win_rate = metrics.get("win_rate", perf.win_rate)
        perf.volatility = metrics.get("volatility", perf.volatility)
        perf.trade_count = int(metrics.get("trade_count", perf.trade_count))
        perf.last_updated = datetime.now()

        return True

    def get_performance(self, strategy_id: Optional[str] = None) -> Dict[str, Any]:
        """Get performance metrics.

        Args:
            strategy_id: Specific strategy or None for all

        Returns:
            Performance metrics dictionary
        """
        if strategy_id:
            perf = self._performance.get(strategy_id)
            if perf:
                return {
                    "strategy_id": perf.strategy_id,
                    "total_return": perf.total_return,
                    "sharpe_ratio": perf.sharpe_ratio,
                    "max_drawdown": perf.max_drawdown,
                    "win_rate": perf.win_rate,
                    "volatility": perf.volatility,
                    "trade_count": perf.trade_count,
                    "last_updated": perf.last_updated,
                }
            return {}

        return {
            sid: {
                "strategy_id": p.strategy_id,
                "total_return": p.total_return,
                "sharpe_ratio": p.sharpe_ratio,
                "max_drawdown": p.max_drawdown,
                "win_rate": p.win_rate,
                "volatility": p.volatility,
            }
            for sid, p in self._performance.items()
        }

    def update_returns(self, returns: pd.DataFrame) -> None:
        """Update returns history for weight calculations.

        Args:
            returns: DataFrame with strategy returns (columns = strategy_ids)
        """
        self._returns_history = returns.copy()

    # ==================== State Management ====================

    def get_state(self) -> PortfolioState:
        """Get current portfolio state.

        Returns:
            PortfolioState object
        """
        self._state.strategies = self._strategies.copy()
        self._state.current_weights = self._current_weights.copy()
        self._state.target_weights = self._target_weights.copy()
        return self._state

    def set_state(self, state: PortfolioState) -> None:
        """Restore portfolio state.

        Args:
            state: PortfolioState to restore
        """
        self._state = state
        self._strategies = state.strategies.copy()
        self._current_weights = state.current_weights.copy()
        self._target_weights = state.target_weights.copy()

    def reset(self) -> None:
        """Reset portfolio to initial state."""
        self._strategies.clear()
        self._strategy_instances.clear()
        self._performance.clear()
        self._current_weights.clear()
        self._target_weights.clear()
        self._returns_history = pd.DataFrame()
        self._rebalancer.reset()

        self._state = PortfolioState(
            portfolio_id=self.portfolio_id,
            total_value=self.initial_capital,
            cash=self.initial_capital,
        )

        self.logger.info("Portfolio reset to initial state")

    def to_dict(self) -> Dict[str, Any]:
        """Convert portfolio to dictionary representation.

        Returns:
            Dictionary with portfolio configuration and state
        """
        return {
            "portfolio_id": self.portfolio_id,
            "weight_method": self.weight_method.value,
            "aggregation_method": self.aggregation_method.value,
            "initial_capital": self.initial_capital,
            "strategies": {
                sid: {
                    "strategy_id": c.strategy_id,
                    "weight": c.weight,
                    "status": c.status.value,
                    "params": c.params,
                }
                for sid, c in self._strategies.items()
            },
            "current_weights": self._current_weights,
            "target_weights": self._target_weights,
            "performance": self.get_performance(),
            "state": {
                "total_value": self._state.total_value,
                "cash": self._state.cash,
                "last_rebalance": self._state.last_rebalance,
            },
        }
