"""
Strategy implementations module.

This module contains:
- Built-in strategy implementations
- Strategy registry for discovering available strategies
- Strategy parameter validation

All strategies take a DataFrame with 'close' column (indexed by date)
and return a Series of signals: 1 = long, 0 = flat.
"""

from typing import Any

import pandas as pd

from jbqlab.types import StrategyFunc, StrategyInfo, StrategyParam

# =============================================================================
# Strategy Registry
# =============================================================================

_STRATEGIES: dict[str, tuple[StrategyFunc, StrategyInfo]] = {}


def register_strategy(info: StrategyInfo):
    """Decorator to register a strategy function.

    Args:
        info: StrategyInfo describing the strategy.

    Returns:
        Decorator function.
    """

    def decorator(func: StrategyFunc) -> StrategyFunc:
        _STRATEGIES[info.name] = (func, info)
        return func

    return decorator


def get_strategy(name: str) -> StrategyFunc:
    """Get a strategy function by name.

    Args:
        name: Name of the strategy.

    Returns:
        Strategy function.

    Raises:
        ValueError: If strategy is not found.
    """
    if name not in _STRATEGIES:
        available = ", ".join(_STRATEGIES.keys())
        raise ValueError(f"Unknown strategy: '{name}'. Available: {available}")
    return _STRATEGIES[name][0]


def get_strategy_info(name: str) -> StrategyInfo:
    """Get strategy metadata by name.

    Args:
        name: Name of the strategy.

    Returns:
        StrategyInfo object.

    Raises:
        ValueError: If strategy is not found.
    """
    if name not in _STRATEGIES:
        available = ", ".join(_STRATEGIES.keys())
        raise ValueError(f"Unknown strategy: '{name}'. Available: {available}")
    return _STRATEGIES[name][1]


def get_available_strategies() -> list[StrategyInfo]:
    """Get list of all registered strategies.

    Returns:
        List of StrategyInfo objects.
    """
    return [info for _, info in _STRATEGIES.values()]


def validate_strategy_params(name: str, params: dict[str, Any]) -> dict[str, Any]:
    """Validate and fill in default values for strategy parameters.

    Args:
        name: Name of the strategy.
        params: User-provided parameters.

    Returns:
        Complete parameter dictionary with defaults filled in.

    Raises:
        ValueError: If required parameters are missing or invalid.
    """
    info = get_strategy_info(name)
    result: dict[str, Any] = {}

    for param in info.params:
        if param.name in params:
            # Validate type
            value = params[param.name]
            try:
                result[param.name] = param.param_type(value)
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f"Invalid value for parameter '{param.name}': {value}. "
                    f"Expected {param.param_type.__name__}. Error: {e}"
                ) from e
        elif param.required:
            raise ValueError(f"Missing required parameter: '{param.name}'")
        else:
            result[param.name] = param.default

    return result


# =============================================================================
# Built-in Strategies
# =============================================================================


@register_strategy(
    StrategyInfo(
        name="buy_and_hold",
        description="Always fully invested (long). Baseline strategy.",
        params=(),
    )
)
def buy_and_hold(df: pd.DataFrame, **_kwargs: Any) -> pd.Series:
    """Buy and hold strategy - always long.

    Args:
        df: DataFrame with 'close' column, indexed by date.
        **kwargs: Ignored (no parameters).

    Returns:
        Series of 1s (always long).
    """
    return pd.Series(1, index=df.index, name="signal")


@register_strategy(
    StrategyInfo(
        name="sma_crossover",
        description=(
            "Go long when fast SMA crosses above slow SMA. "
            "Flat when fast SMA is below slow SMA."
        ),
        params=(
            StrategyParam(
                name="fast",
                description="Fast moving average window (days)",
                param_type=int,
                default=10,
            ),
            StrategyParam(
                name="slow",
                description="Slow moving average window (days)",
                param_type=int,
                default=30,
            ),
        ),
    )
)
def sma_crossover(df: pd.DataFrame, fast: int = 10, slow: int = 30, **_kwargs: Any) -> pd.Series:
    """SMA crossover strategy.

    Go long when the fast SMA is above the slow SMA.

    Args:
        df: DataFrame with 'close' column, indexed by date.
        fast: Fast moving average window (default 10).
        slow: Slow moving average window (default 30).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).

    Raises:
        ValueError: If parameters are invalid.
    """
    if fast <= 0 or slow <= 0:
        raise ValueError("SMA windows must be positive integers")
    if fast >= slow:
        raise ValueError(f"fast ({fast}) must be less than slow ({slow})")

    close = df["close"]
    sma_fast = close.rolling(window=fast, min_periods=fast).mean()
    sma_slow = close.rolling(window=slow, min_periods=slow).mean()

    # Signal: 1 when fast > slow, 0 otherwise
    signal = (sma_fast > sma_slow).astype(int)

    # Fill NaN with 0 (no position during warmup)
    signal = signal.fillna(0).astype(int)
    signal.name = "signal"

    return signal


@register_strategy(
    StrategyInfo(
        name="mean_reversion",
        description=(
            "Go long when price is z_entry standard deviations below the rolling mean. "
            "Exit when price returns to the mean. Long-only."
        ),
        params=(
            StrategyParam(
                name="window",
                description="Lookback window for mean and std calculation (days)",
                param_type=int,
                default=20,
            ),
            StrategyParam(
                name="z_entry",
                description="Z-score threshold to enter long (e.g., -2.0 means 2 std below mean)",
                param_type=float,
                default=-2.0,
            ),
        ),
    )
)
def mean_reversion(
    df: pd.DataFrame, window: int = 20, z_entry: float = -2.0, **_kwargs: Any
) -> pd.Series:
    """Mean reversion strategy (long-only).

    Enter long when price falls z_entry standard deviations below the rolling mean.
    Exit when price returns to or above the rolling mean.

    Args:
        df: DataFrame with 'close' column, indexed by date.
        window: Lookback window for rolling statistics (default 20).
        z_entry: Z-score threshold to enter long (default -2.0, i.e., 2 std below).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).

    Raises:
        ValueError: If parameters are invalid.
    """
    if window <= 1:
        raise ValueError("window must be at least 2")
    if z_entry >= 0:
        raise ValueError("z_entry should be negative for long-only mean reversion")

    close = df["close"]
    rolling_mean = close.rolling(window=window, min_periods=window).mean()
    rolling_std = close.rolling(window=window, min_periods=window).std()

    # Calculate z-score
    z_score = (close - rolling_mean) / rolling_std

    # State machine for signals
    # Enter long when z_score < z_entry
    # Exit when z_score >= 0 (price returned to mean)
    signal = pd.Series(0, index=df.index, name="signal")

    position = 0  # Current position state
    for i in range(len(df)):
        z = z_score.iloc[i]

        if pd.isna(z):
            signal.iloc[i] = 0
            continue

        if position == 0:  # Not in position
            if z <= z_entry:
                position = 1  # Enter long
        else:  # In position
            if z >= 0:
                position = 0  # Exit

        signal.iloc[i] = position

    return signal


@register_strategy(
    StrategyInfo(
        name="momentum",
        description=(
            "Go long when the return over the lookback period is positive. "
            "Uses rate of change (ROC) as momentum indicator."
        ),
        params=(
            StrategyParam(
                name="lookback",
                description="Lookback period for momentum calculation (days)",
                param_type=int,
                default=20,
            ),
            StrategyParam(
                name="threshold",
                description="Minimum ROC threshold to go long (e.g., 0.0 means any positive momentum)",
                param_type=float,
                default=0.0,
            ),
        ),
    )
)
def momentum(
    df: pd.DataFrame, lookback: int = 20, threshold: float = 0.0, **_kwargs: Any
) -> pd.Series:
    """Momentum strategy based on Rate of Change (ROC).

    Go long when the price change over the lookback period exceeds the threshold.

    Args:
        df: DataFrame with 'close' column, indexed by date.
        lookback: Lookback period for ROC calculation (default 20).
        threshold: Minimum ROC to trigger long signal (default 0.0).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).

    Raises:
        ValueError: If parameters are invalid.
    """
    if lookback <= 0:
        raise ValueError("lookback must be positive")

    close = df["close"]

    # Rate of Change: (current - past) / past
    roc = (close - close.shift(lookback)) / close.shift(lookback)

    # Signal: 1 when ROC > threshold, 0 otherwise
    signal = (roc > threshold).astype(int)

    # Fill NaN with 0 (no position during warmup)
    signal = signal.fillna(0).astype(int)
    signal.name = "signal"

    return signal


@register_strategy(
    StrategyInfo(
        name="rsi",
        description=(
            "RSI (Relative Strength Index) strategy. "
            "Go long when RSI falls below oversold level, exit when RSI rises above overbought level."
        ),
        params=(
            StrategyParam(
                name="period",
                description="RSI calculation period (days)",
                param_type=int,
                default=14,
            ),
            StrategyParam(
                name="oversold",
                description="RSI level considered oversold (buy signal)",
                param_type=float,
                default=30.0,
            ),
            StrategyParam(
                name="overbought",
                description="RSI level considered overbought (sell signal)",
                param_type=float,
                default=70.0,
            ),
        ),
    )
)
def rsi_strategy(
    df: pd.DataFrame,
    period: int = 14,
    oversold: float = 30.0,
    overbought: float = 70.0,
    **_kwargs: Any,
) -> pd.Series:
    """RSI-based mean reversion strategy.

    Enter long when RSI drops below oversold threshold.
    Exit when RSI rises above overbought threshold.

    Args:
        df: DataFrame with 'close' column, indexed by date.
        period: RSI calculation period (default 14).
        oversold: RSI threshold to enter long (default 30).
        overbought: RSI threshold to exit long (default 70).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).

    Raises:
        ValueError: If parameters are invalid.
    """
    if period <= 0:
        raise ValueError("period must be positive")
    if not 0 <= oversold < overbought <= 100:
        raise ValueError("Must have 0 <= oversold < overbought <= 100")

    close = df["close"]

    # Calculate RSI
    delta = close.diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)

    avg_gain = gain.rolling(window=period, min_periods=period).mean()
    avg_loss = loss.rolling(window=period, min_periods=period).mean()

    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))

    # State machine for signals
    signal = pd.Series(0, index=df.index, name="signal")

    position = 0
    for i in range(len(df)):
        r = rsi.iloc[i]

        if pd.isna(r):
            signal.iloc[i] = 0
            continue

        if position == 0:  # Not in position
            if r <= oversold:
                position = 1  # Enter long
        else:  # In position
            if r >= overbought:
                position = 0  # Exit

        signal.iloc[i] = position

    return signal


@register_strategy(
    StrategyInfo(
        name="bollinger_bands",
        description=(
            "Go long when price touches or crosses below the lower Bollinger Band. "
            "Exit when price touches or crosses above the upper band. Mean reversion strategy."
        ),
        params=(
            StrategyParam(
                name="window",
                description="Lookback window for moving average and std calculation",
                param_type=int,
                default=20,
            ),
            StrategyParam(
                name="num_std",
                description="Number of standard deviations for band width",
                param_type=float,
                default=2.0,
            ),
        ),
    )
)
def bollinger_bands(
    df: pd.DataFrame,
    window: int = 20,
    num_std: float = 2.0,
    **_kwargs: Any,
) -> pd.Series:
    """Bollinger Bands mean reversion strategy.

    Enter long when price touches or goes below lower band.
    Exit when price touches or goes above upper band.

    Args:
        df: DataFrame with 'close' column, indexed by date.
        window: Lookback window for SMA and standard deviation (default 20).
        num_std: Number of standard deviations for bands (default 2.0).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).

    Raises:
        ValueError: If parameters are invalid.
    """
    if window <= 1:
        raise ValueError("window must be at least 2")
    if num_std <= 0:
        raise ValueError("num_std must be positive")

    close = df["close"]

    # Calculate Bollinger Bands
    sma = close.rolling(window=window, min_periods=window).mean()
    std = close.rolling(window=window, min_periods=window).std()

    upper_band = sma + (num_std * std)
    lower_band = sma - (num_std * std)

    # State machine for signals
    signal = pd.Series(0, index=df.index, name="signal")

    position = 0
    for i in range(len(df)):
        price = close.iloc[i]
        lower = lower_band.iloc[i]
        upper = upper_band.iloc[i]

        if pd.isna(lower) or pd.isna(upper):
            signal.iloc[i] = 0
            continue

        if position == 0:  # Not in position
            if price <= lower:
                position = 1  # Enter long at lower band
        else:  # In position
            if price >= upper:
                position = 0  # Exit at upper band

        signal.iloc[i] = position

    return signal


@register_strategy(
    StrategyInfo(
        name="dual_momentum",
        description=(
            "Dual momentum strategy combining absolute and relative momentum. "
            "Go long only when both absolute momentum (vs risk-free) and "
            "price momentum are positive."
        ),
        params=(
            StrategyParam(
                name="lookback",
                description="Lookback period for momentum calculation",
                param_type=int,
                default=12,  # 12 months if using monthly data, 12 days if daily
            ),
        ),
    )
)
def dual_momentum(
    df: pd.DataFrame,
    lookback: int = 12,
    **_kwargs: Any,
) -> pd.Series:
    """Dual momentum strategy.

    Combines absolute momentum (is the return positive?) with
    relative momentum (is it trending up?).

    Args:
        df: DataFrame with 'close' column, indexed by date.
        lookback: Lookback period for momentum (default 12).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).
    """
    if lookback <= 0:
        raise ValueError("lookback must be positive")

    close = df["close"]

    # Calculate momentum (return over lookback period)
    momentum_return = (close - close.shift(lookback)) / close.shift(lookback)

    # Absolute momentum: is the return positive?
    absolute_mom = momentum_return > 0

    # Relative momentum: is the short-term trend positive?
    # Use half the lookback for short-term
    short_lookback = max(1, lookback // 2)
    short_momentum = (close - close.shift(short_lookback)) / close.shift(short_lookback)
    relative_mom = short_momentum > 0

    # Go long only when both are positive
    signal = (absolute_mom & relative_mom).astype(int)
    signal = signal.fillna(0).astype(int)
    signal.name = "signal"

    return signal


@register_strategy(
    StrategyInfo(
        name="macd",
        description=(
            "MACD (Moving Average Convergence Divergence) strategy. "
            "Go long when MACD line crosses above signal line. "
            "Exit when MACD crosses below signal line."
        ),
        params=(
            StrategyParam(
                name="fast_period",
                description="Fast EMA period",
                param_type=int,
                default=12,
            ),
            StrategyParam(
                name="slow_period",
                description="Slow EMA period",
                param_type=int,
                default=26,
            ),
            StrategyParam(
                name="signal_period",
                description="Signal line EMA period",
                param_type=int,
                default=9,
            ),
        ),
    )
)
def macd(
    df: pd.DataFrame,
    fast_period: int = 12,
    slow_period: int = 26,
    signal_period: int = 9,
    **_kwargs: Any,
) -> pd.Series:
    """MACD crossover strategy.

    Go long when MACD crosses above signal line.
    Exit when MACD crosses below signal line.

    Args:
        df: DataFrame with 'close' column, indexed by date.
        fast_period: Fast EMA period (default 12).
        slow_period: Slow EMA period (default 26).
        signal_period: Signal line EMA period (default 9).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).

    Raises:
        ValueError: If parameters are invalid.
    """
    if fast_period <= 0:
        raise ValueError("fast_period must be positive")
    if slow_period <= 0:
        raise ValueError("slow_period must be positive")
    if signal_period <= 0:
        raise ValueError("signal_period must be positive")
    if fast_period >= slow_period:
        raise ValueError("fast_period must be less than slow_period")

    close = df["close"]

    # Calculate EMAs
    fast_ema = close.ewm(span=fast_period, adjust=False).mean()
    slow_ema = close.ewm(span=slow_period, adjust=False).mean()

    # MACD line = fast EMA - slow EMA
    macd_line = fast_ema - slow_ema

    # Signal line = EMA of MACD line
    signal_line = macd_line.ewm(span=signal_period, adjust=False).mean()

    # Go long when MACD > Signal, flat otherwise
    signal = (macd_line > signal_line).astype(int)
    signal = signal.fillna(0).astype(int)
    signal.name = "signal"

    return signal


@register_strategy(
    StrategyInfo(
        name="donchian_channel",
        description=(
            "Donchian Channel breakout strategy (trend-following). "
            "Go long when price breaks above the upper channel. "
            "Exit when price breaks below the lower channel."
        ),
        params=(
            StrategyParam(
                name="entry_period",
                description="Lookback period for entry channel",
                param_type=int,
                default=20,
            ),
            StrategyParam(
                name="exit_period",
                description="Lookback period for exit channel",
                param_type=int,
                default=10,
            ),
        ),
    )
)
def donchian_channel(
    df: pd.DataFrame,
    entry_period: int = 20,
    exit_period: int = 10,
    **_kwargs: Any,
) -> pd.Series:
    """Donchian Channel breakout strategy.

    Enter long when price breaks above the entry period high.
    Exit when price breaks below the exit period low.

    Args:
        df: DataFrame with 'close' (and optionally 'high', 'low') columns.
        entry_period: Lookback for entry channel (default 20).
        exit_period: Lookback for exit channel (default 10).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).
    """
    if entry_period <= 0:
        raise ValueError("entry_period must be positive")
    if exit_period <= 0:
        raise ValueError("exit_period must be positive")

    # Use high/low if available, else use close
    high = df["high"] if "high" in df.columns else df["close"]
    low = df["low"] if "low" in df.columns else df["close"]
    close = df["close"]

    # Calculate channels (shifted by 1 to avoid look-ahead bias)
    upper_channel = high.rolling(window=entry_period).max().shift(1)
    lower_channel = low.rolling(window=exit_period).min().shift(1)

    # State machine for signals
    signal = pd.Series(0, index=df.index, name="signal")

    position = 0
    for i in range(len(df)):
        price = close.iloc[i]
        upper = upper_channel.iloc[i]
        lower = lower_channel.iloc[i]

        if pd.isna(upper) or pd.isna(lower):
            signal.iloc[i] = 0
            continue

        if position == 0:  # Not in position
            if price > upper:
                position = 1  # Breakout entry
        else:  # In position
            if price < lower:
                position = 0  # Breakdown exit

        signal.iloc[i] = position

    return signal


@register_strategy(
    StrategyInfo(
        name="volatility_breakout",
        description=(
            "Volatility breakout strategy based on ATR (Average True Range). "
            "Enter long when price breaks above previous close + multiplier * ATR. "
            "Exit at end of day or when price drops below entry - ATR."
        ),
        params=(
            StrategyParam(
                name="atr_period",
                description="Period for ATR calculation",
                param_type=int,
                default=14,
            ),
            StrategyParam(
                name="multiplier",
                description="ATR multiplier for breakout threshold",
                param_type=float,
                default=0.5,
            ),
        ),
    )
)
def volatility_breakout(
    df: pd.DataFrame,
    atr_period: int = 14,
    multiplier: float = 0.5,
    **_kwargs: Any,
) -> pd.Series:
    """Volatility breakout strategy using ATR.

    Enter when price exceeds previous close by a multiple of ATR.
    This is a simplified daily-based version.

    Args:
        df: DataFrame with 'close' column (and optionally 'high', 'low').
        atr_period: Period for ATR calculation (default 14).
        multiplier: ATR multiplier for entry threshold (default 0.5).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).
    """
    if atr_period <= 0:
        raise ValueError("atr_period must be positive")
    if multiplier <= 0:
        raise ValueError("multiplier must be positive")

    close = df["close"]

    # Calculate True Range components
    if "high" in df.columns and "low" in df.columns:
        high = df["high"]
        low = df["low"]
        prev_close = close.shift(1)

        tr1 = high - low
        tr2 = (high - prev_close).abs()
        tr3 = (low - prev_close).abs()

        true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
    else:
        # Fallback: use absolute daily returns as proxy for range
        true_range = close.diff().abs()

    # Average True Range
    atr = true_range.rolling(window=atr_period).mean()

    # Breakout threshold
    prev_close = close.shift(1)
    breakout_level = prev_close + (multiplier * atr.shift(1))

    # Signal: enter if price exceeds breakout level
    signal = (close > breakout_level).astype(int)
    signal = signal.fillna(0).astype(int)
    signal.name = "signal"

    return signal


@register_strategy(
    StrategyInfo(
        name="trailing_stop",
        description=(
            "Always long with a trailing stop strategy. "
            "Position is held with a trailing stop at a percentage below the highest price. "
            "Re-enter after stop-out when price makes new high."
        ),
        params=(
            StrategyParam(
                name="trail_pct",
                description="Trailing stop percentage below high",
                param_type=float,
                default=0.05,  # 5%
            ),
        ),
    )
)
def trailing_stop(
    df: pd.DataFrame,
    trail_pct: float = 0.05,
    **_kwargs: Any,
) -> pd.Series:
    """Trailing stop strategy.

    Stay long as long as price doesn't fall below trailing stop.
    Re-enter when price makes new high.

    Args:
        df: DataFrame with 'close' column.
        trail_pct: Trailing stop percentage (default 0.05 = 5%).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).
    """
    if trail_pct <= 0 or trail_pct >= 1:
        raise ValueError("trail_pct must be between 0 and 1")

    close = df["close"]
    signal = pd.Series(0, index=df.index, name="signal")

    position = 0
    highest = close.iloc[0]

    for i in range(len(df)):
        price = close.iloc[i]

        if position == 0:  # Flat
            # Enter on new high
            if price > highest:
                position = 1
                highest = price
        else:  # In position
            # Update trailing stop
            if price > highest:
                highest = price

            # Check stop
            stop_level = highest * (1 - trail_pct)
            if price < stop_level:
                position = 0
                highest = price  # Reset for next entry

        signal.iloc[i] = position

    return signal


@register_strategy(
    StrategyInfo(
        name="triple_sma",
        description=(
            "Triple SMA filter strategy. Go long only when "
            "short SMA > medium SMA > long SMA (strong uptrend)."
        ),
        params=(
            StrategyParam(
                name="short_window",
                description="Short SMA period",
                param_type=int,
                default=10,
            ),
            StrategyParam(
                name="medium_window",
                description="Medium SMA period",
                param_type=int,
                default=20,
            ),
            StrategyParam(
                name="long_window",
                description="Long SMA period",
                param_type=int,
                default=50,
            ),
        ),
    )
)
def triple_sma(
    df: pd.DataFrame,
    short_window: int = 10,
    medium_window: int = 20,
    long_window: int = 50,
    **_kwargs: Any,
) -> pd.Series:
    """Triple SMA trend filter strategy.

    Long when short > medium > long SMA (aligned uptrend).

    Args:
        df: DataFrame with 'close' column.
        short_window: Short SMA period (default 10).
        medium_window: Medium SMA period (default 20).
        long_window: Long SMA period (default 50).
        **kwargs: Ignored.

    Returns:
        Series of signals (1 = long, 0 = flat).
    """
    if short_window <= 0:
        raise ValueError("short_window must be positive")
    if medium_window <= 0:
        raise ValueError("medium_window must be positive")
    if long_window <= 0:
        raise ValueError("long_window must be positive")
    if not (short_window < medium_window < long_window):
        raise ValueError("Must have short_window < medium_window < long_window")

    close = df["close"]

    short_sma = close.rolling(window=short_window, min_periods=short_window).mean()
    medium_sma = close.rolling(window=medium_window, min_periods=medium_window).mean()
    long_sma = close.rolling(window=long_window, min_periods=long_window).mean()

    # Aligned uptrend condition
    uptrend = (short_sma > medium_sma) & (medium_sma > long_sma)

    signal = uptrend.astype(int)
    signal = signal.fillna(0).astype(int)
    signal.name = "signal"

    return signal
