"""Statistical Arbitrage (Pairs Trading) Strategy."""

from decimal import Decimal
from typing import Dict, Any, Optional, List
import pandas as pd
import numpy as np

from ..base import BaseStrategy, StrategySignal, SignalType
from src.data.pricing_service import get_pricing_service
from src.core.logger import get_logger

logger = get_logger(__name__)

class PairsTradingStrategy(BaseStrategy):
    """
    Statistical Arbitrage Strategy that trades the spread between two correlated assets.
    
    Logic:
    1. Calculate Spread = Price(A) / Price(B)
    2. Calculate Z-Score of the spread over a rolling window.
    3. Mean Reversion signals:
       - Z-Score > Threshold: Short Spread (Sell A, Buy B)
       - Z-Score < -Threshold: Long Spread (Buy A, Sell B)
       - Z-Score approx 0: Close/Neutralize (Exit both)
    """

    def __init__(self, name: str, parameters: Dict[str, Any], timeframes: List[str] = None):
        super().__init__(name, parameters, timeframes)
        self.second_symbol = parameters.get('second_symbol')
        self.lookback_window = int(parameters.get('lookback_period', 20))
        self.z_threshold = float(parameters.get('z_score_threshold', 2.0))
        self.pricing_service = get_pricing_service()
        
        if not self.second_symbol:
            logger.warning(f"PairsTradingStrategy {name} initialized without 'second_symbol'")

    async def on_tick(self, symbol: str, price: Decimal, timeframe: str, data: Dict[str, Any]) -> Optional[StrategySignal]:
        """Check for pairs trading signals."""
        if not self.second_symbol or not self.enabled:
            return None

        # We only process on the "Primary" symbol's tick to avoid double-processing
        # (Assuming the strategy is assigned to the primary symbol in the DB)
        
        # 1. Fetch data for BOTH symbols
        # We need historical data to calculate Z-Score
        try:
            # Fetch for Primary (Symbol A)
            ohlcv_a = self.pricing_service.get_ohlcv(
                symbol=symbol,
                timeframe=timeframe,
                limit=self.lookback_window + 5
            )
            
            # Fetch for Secondary (Symbol B)
            ohlcv_b = self.pricing_service.get_ohlcv(
                symbol=self.second_symbol,
                timeframe=timeframe,
                limit=self.lookback_window + 5
            )
            
            if not ohlcv_a or not ohlcv_b:
                return None
                
            # Convert to Series for pandas calc
            # OHLCV format: [timestamp, open, high, low, close, volume]
            closes_a = pd.Series([float(c[4]) for c in ohlcv_a])
            closes_b = pd.Series([float(c[4]) for c in ohlcv_b])
            
            # Ensure equal length if fetched data differs slightly
            min_len = min(len(closes_a), len(closes_b))
            closes_a = closes_a.iloc[-min_len:]
            closes_b = closes_b.iloc[-min_len:]

            # 2. Calculate Spread and Z-Score
            # Spread = A / B
            spread = closes_a / closes_b
            
            # Rolling statistics
            rolling_mean = spread.rolling(window=self.lookback_window).mean()
            rolling_std = spread.rolling(window=self.lookback_window).std()
            
            current_spread = spread.iloc[-1]
            current_mean = rolling_mean.iloc[-1]
            current_std = rolling_std.iloc[-1]
            
            if pd.isna(current_std) or current_std == 0:
                return None
                
            z_score = (current_spread - current_mean) / current_std
            
            self.logger.info(
                f"Pairs {symbol}/{self.second_symbol}: Spread={current_spread:.4f}, Z-Score={z_score:.2f}"
            )

            # 3. Generate Signals
            # Strategy:
            # If Z > Threshold -> Spread is too high (A is expensive, B is cheap) -> Sell A, Buy B
            # If Z < -Threshold -> Spread is too low (A is cheap, B is expensive) -> Buy A, Sell B
            # If abs(Z) < 0.5 -> Mean reverted -> Close Positions (optional, or just go to neutral)
            
            signal_type = None
            primary_side = "hold"
            secondary_side = "hold"
            
            if z_score > self.z_threshold:
                # Sell A, Buy B
                signal_type = SignalType.SELL
                primary_side = "sell"
                secondary_side = "buy"
                
            elif z_score < -self.z_threshold:
                # Buy A, Sell B
                signal_type = SignalType.BUY
                primary_side = "buy"
                secondary_side = "sell"
                
            # TODO: Logic for closing when Z ~ 0 can be added here
            
            if signal_type:
                # Create Signal for Primary
                signal = StrategySignal(
                    signal_type=signal_type,
                    symbol=symbol,
                    strength=min(abs(z_score) / self.z_threshold, 1.0), # Strength capped at 1.0
                    price=price,
                    metadata={
                        "strategy": "pairs_trading",
                        "z_score": float(z_score),
                        "spread": float(current_spread),
                        "secondary_symbol": self.second_symbol,
                        "secondary_action": secondary_side,
                        "description": f"Pairs Arbitrage: Z-Score {z_score:.2f} triggered {primary_side} {symbol} / {secondary_side} {self.second_symbol}"
                    }
                )
                return signal

        except Exception as e:
            self.logger.error(f"Error in PairsTradingStrategy: {e}")
            return None
            
        return None

    def on_signal(self, signal: StrategySignal) -> Optional[StrategySignal]:
        """Process signal (pass-through)."""
        return signal
