"""Stop-loss logic."""

from decimal import Decimal
from typing import Dict, Optional, Any
import pandas as pd
from src.core.logger import get_logger
from src.data.indicators import get_indicators

logger = get_logger(__name__)


class StopLossManager:
    """Manages stop-loss orders."""
    
    def __init__(self):
        """Initialize stop-loss manager."""
        self.stop_losses: Dict[int, Dict[str, Any]] = {}  # position_id -> stop config
        self.logger = get_logger(__name__)
        self.indicators = get_indicators()
    
    def set_stop_loss(
        self,
        position_id: int,
        stop_price: Decimal,
        trailing: bool = False,
        trail_percent: Optional[Decimal] = None,
        use_atr: bool = False,
        atr_multiplier: Decimal = Decimal('2.0'),
        atr_period: int = 14,
        ohlcv_data: Optional[pd.DataFrame] = None
    ):
        """Set stop-loss for position.
        
        Args:
            position_id: Position ID
            stop_price: Stop price (ignored if use_atr=True)
            trailing: Enable trailing stop
            trail_percent: Trail percentage (ignored if use_atr=True)
            use_atr: Use ATR-based stop loss
            atr_multiplier: ATR multiplier for stop distance (default 2.0)
            atr_period: ATR period (default 14)
            ohlcv_data: OHLCV DataFrame for ATR calculation
        """
        if use_atr and ohlcv_data is not None and len(ohlcv_data) >= atr_period:
            # Calculate ATR-based stop
            try:
                high = ohlcv_data['high']
                low = ohlcv_data['low']
                close = ohlcv_data['close']
                
                atr = self.indicators.atr(high, low, close, period=atr_period)
                current_atr = Decimal(str(atr.iloc[-1])) if not pd.isna(atr.iloc[-1]) else None
                
                if current_atr is not None:
                    current_price = Decimal(str(close.iloc[-1]))
                    atr_distance = current_atr * atr_multiplier
                    
                    # For long positions, stop is below entry
                    # For short positions, stop is above entry
                    # We'll determine direction from stop_price vs current_price
                    if stop_price < current_price:
                        # Long position - stop below
                        stop_price = current_price - atr_distance
                    else:
                        # Short position - stop above
                        stop_price = current_price + atr_distance
                    
                    self.logger.info(
                        f"Set ATR-based stop-loss for position {position_id}: "
                        f"ATR={current_atr}, multiplier={atr_multiplier}, "
                        f"stop_price={stop_price}"
                    )
            except Exception as e:
                self.logger.warning(f"Failed to calculate ATR-based stop, using provided stop_price: {e}")
        
        self.stop_losses[position_id] = {
            'stop_price': stop_price,
            'trailing': trailing,
            'trail_percent': trail_percent,
            'use_atr': use_atr,
            'atr_multiplier': atr_multiplier if use_atr else None,
            'atr_period': atr_period if use_atr else None,
            'highest_price': stop_price,  # For trailing stops
        }
        if not use_atr:
            self.logger.info(f"Set stop-loss for position {position_id} at {stop_price}")
    
    def check_stop_loss(
        self,
        position_id: int,
        current_price: Decimal,
        is_long: bool = True,
        ohlcv_data: Optional[pd.DataFrame] = None
    ) -> bool:
        """Check if stop-loss should trigger.
        
        Args:
            position_id: Position ID
            current_price: Current market price
            is_long: True for long position
            ohlcv_data: OHLCV DataFrame for ATR-based trailing stops
            
        Returns:
            True if stop-loss should trigger
        """
        if position_id not in self.stop_losses:
            return False
        
        stop_config = self.stop_losses[position_id]
        stop_price = stop_config['stop_price']
        use_atr = stop_config.get('use_atr', False)
        
        if stop_config['trailing']:
            # Update trailing stop
            if use_atr and ohlcv_data is not None and stop_config.get('atr_period'):
                # ATR-based trailing stop
                try:
                    atr_period = stop_config['atr_period']
                    atr_multiplier = stop_config.get('atr_multiplier', Decimal('2.0'))
                    
                    if len(ohlcv_data) >= atr_period:
                        high = ohlcv_data['high']
                        low = ohlcv_data['low']
                        close = ohlcv_data['close']
                        
                        atr = self.indicators.atr(high, low, close, period=atr_period)
                        current_atr = Decimal(str(atr.iloc[-1])) if not pd.isna(atr.iloc[-1]) else None
                        
                        if current_atr is not None:
                            atr_distance = current_atr * atr_multiplier
                            
                            if is_long:
                                # Long position: trailing stop moves up
                                if current_price > stop_config['highest_price']:
                                    stop_config['highest_price'] = current_price
                                    stop_price = current_price - atr_distance
                                    stop_config['stop_price'] = stop_price
                            else:
                                # Short position: trailing stop moves down
                                if current_price < stop_config['highest_price']:
                                    stop_config['highest_price'] = current_price
                                    stop_price = current_price + atr_distance
                                    stop_config['stop_price'] = stop_price
                except Exception as e:
                    self.logger.warning(f"Error updating ATR-based trailing stop: {e}")
            else:
                # Percentage-based trailing stop
                if is_long:
                    if current_price > stop_config['highest_price']:
                        stop_config['highest_price'] = current_price
                        if stop_config.get('trail_percent'):
                            stop_price = current_price * (1 - stop_config['trail_percent'])
                            stop_config['stop_price'] = stop_price
                else:
                    if current_price < stop_config['highest_price']:
                        stop_config['highest_price'] = current_price
                        if stop_config.get('trail_percent'):
                            stop_price = current_price * (1 + stop_config['trail_percent'])
                            stop_config['stop_price'] = stop_price
        
        # Check trigger
        if is_long:
            return current_price <= stop_price
        else:
            return current_price >= stop_price
    
    def calculate_atr_stop(
        self,
        entry_price: Decimal,
        is_long: bool,
        ohlcv_data: pd.DataFrame,
        atr_multiplier: Decimal = Decimal('2.0'),
        atr_period: int = 14
    ) -> Decimal:
        """Calculate ATR-based stop loss price.
        
        Args:
            entry_price: Entry price
            is_long: True for long position
            ohlcv_data: OHLCV DataFrame
            atr_multiplier: ATR multiplier (default 2.0)
            atr_period: ATR period (default 14)
            
        Returns:
            Stop loss price
        """
        if len(ohlcv_data) < atr_period:
            # Fallback to 2% stop if insufficient data
            if is_long:
                return entry_price * Decimal('0.98')
            else:
                return entry_price * Decimal('1.02')
        
        try:
            high = ohlcv_data['high']
            low = ohlcv_data['low']
            close = ohlcv_data['close']
            
            atr = self.indicators.atr(high, low, close, period=atr_period)
            current_atr = Decimal(str(atr.iloc[-1])) if not pd.isna(atr.iloc[-1]) else None
            
            if current_atr is not None:
                atr_distance = current_atr * atr_multiplier
                
                if is_long:
                    # Long position: stop below entry
                    return entry_price - atr_distance
                else:
                    # Short position: stop above entry
                    return entry_price + atr_distance
        except Exception as e:
            self.logger.warning(f"Error calculating ATR stop, using fallback: {e}")
        
        # Fallback to percentage-based stop
        if is_long:
            return entry_price * Decimal('0.98')
        else:
            return entry_price * Decimal('1.02')
    
    def remove_stop_loss(self, position_id: int):
        """Remove stop-loss for position.
        
        Args:
            position_id: Position ID
        """
        if position_id in self.stop_losses:
            del self.stop_losses[position_id]
            self.logger.info(f"Removed stop-loss for position {position_id}")

