from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Optional
from datetime import datetime, timezone

if TYPE_CHECKING:
    from src.stats import TradingStats

logger = logging.getLogger(__name__)

class DrawdownMonitor:
    """
    Tracks portfolio balance in memory to calculate max drawdown in real-time
    and persists its state to the database.
    """

    def __init__(self, stats: "TradingStats"):
        """
        Initializes the DrawdownMonitor.

        Args:
            stats: An instance of the TradingStats class to persist state.
        """
        self.stats = stats
        self.peak_balance = 0.0
        self.peak_balance_timestamp: Optional[datetime] = None
        self.max_drawdown_pct = 0.0
        self.current_drawdown_pct = 0.0
        self._load_state()

    def _load_state(self):
        """Load the persisted state from the database."""
        try:
            peak_balance_str = self.stats._get_metadata('drawdown_peak_balance')
            max_drawdown_pct_str = self.stats._get_metadata('drawdown_max_drawdown_pct')
            peak_balance_timestamp_str = self.stats._get_metadata('drawdown_peak_balance_timestamp')
            
            self.peak_balance = float(peak_balance_str) if peak_balance_str else 0.0
            self.max_drawdown_pct = float(max_drawdown_pct_str) if max_drawdown_pct_str else 0.0
            if peak_balance_timestamp_str:
                self.peak_balance_timestamp = datetime.fromisoformat(peak_balance_timestamp_str)
            
            # If peak balance is zero, initialize it with the initial account balance.
            if self.peak_balance == 0.0:
                initial_balance_str = self.stats._get_metadata('initial_balance')
                if initial_balance_str:
                    self.peak_balance = float(initial_balance_str)
                    self.peak_balance_timestamp = datetime.now(timezone.utc)
            
            logger.info(f"DrawdownMonitor state loaded: Peak Balance=${self.peak_balance:,.2f}, Max Drawdown={self.max_drawdown_pct:.2f}%")

        except Exception as e:
            logger.error(f"Error loading DrawdownMonitor state: {e}", exc_info=True)

    def _save_state(self):
        """Save the current state to the database."""
        try:
            self.stats._set_metadata('drawdown_peak_balance', str(self.peak_balance))
            self.stats._set_metadata('drawdown_max_drawdown_pct', str(self.max_drawdown_pct))
            if self.peak_balance_timestamp:
                self.stats._set_metadata('drawdown_peak_balance_timestamp', self.peak_balance_timestamp.isoformat())
            logger.debug("DrawdownMonitor state saved.")
        except Exception as e:
            logger.error(f"Error saving DrawdownMonitor state: {e}", exc_info=True)

    def update_balance(self, current_balance: float):
        """
        Update the balance and recalculate the drawdown if necessary.

        Args:
            current_balance: The current total balance of the portfolio.
        """
        state_changed = False
        
        if current_balance > self.peak_balance:
            self.peak_balance = current_balance
            self.peak_balance_timestamp = datetime.now(timezone.utc)
            state_changed = True
        
        if self.peak_balance > 0:
            drawdown = (self.peak_balance - current_balance) / self.peak_balance
            self.current_drawdown_pct = drawdown * 100
            
            # Only update if the new drawdown is significantly larger
            if self.current_drawdown_pct > self.max_drawdown_pct + 0.01:
                self.max_drawdown_pct = self.current_drawdown_pct
                state_changed = True
        
        if state_changed:
            self._save_state()

    def get_max_drawdown(self) -> float:
        """Returns the maximum drawdown percentage."""
        return self.max_drawdown_pct

    def get_current_drawdown(self) -> float:
        """Returns the current drawdown percentage since the last peak."""
        return self.current_drawdown_pct

    def get_peak_balance_timestamp(self) -> Optional[datetime]:
        """Returns the timestamp of the last peak balance."""
        return self.peak_balance_timestamp