"""Statistical arbitrage using cointegration-based pairs trading.

More rigorous than spread convergence: uses OLS hedge ratio,
ADF stationarity test, and half-life filtering.
"""

from __future__ import annotations

import logging
import time
from collections import deque
from dataclasses import dataclass
from typing import Callable

from horizon.context import Context

from ._types import StatArbResult

logger = logging.getLogger(__name__)


@dataclass
class StatArbConfig:
    """Configuration for statistical arbitrage."""

    pair: tuple[str, str]
    feeds: tuple[str, str]
    lookback: int = 200
    entry_zscore: float = 2.0
    exit_zscore: float = 0.5
    stop_zscore: float = 4.0
    recalibrate_every: int = 50
    min_half_life: float = 5.0
    max_half_life: float = 200.0


def stat_arb(
    config: StatArbConfig,
    size: float = 10.0,
    auto_execute: bool = False,
    cooldown: float = 30.0,
) -> Callable[[Context], None]:
    """Create a pipeline function for cointegration-based stat arb.

    Collects price histories, periodically recalibrates the hedge ratio
    via OLS cointegration test, checks half-life bounds, and generates
    entry/exit/stop signals based on residual z-score.

    Args:
        config: StatArbConfig with pair, feeds, and parameters.
        size: Trade size per leg.
        auto_execute: If True, execute trades automatically.
        cooldown: Seconds between executions.

    Returns:
        Pipeline function: (Context) -> None
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    prices_a: deque[float] = deque(maxlen=config.lookback)
    prices_b: deque[float] = deque(maxlen=config.lookback)
    hedge_ratio = 1.0
    residuals: list[float] = []
    ticks_since_recalib = 0
    position_state = "flat"  # "flat" | "long_a_short_b" | "long_b_short_a"
    last_exec_time = 0.0
    last_half_life = 0.0
    last_adf = 0.0

    market_a, market_b = config.pair
    feed_a, feed_b = config.feeds

    def _scanner(ctx: Context) -> None:
        nonlocal hedge_ratio, residuals, ticks_since_recalib
        nonlocal position_state, last_exec_time, last_half_life, last_adf

        engine = ctx.params.get("engine")
        if engine is None:
            return None

        snap_a = engine.feed_snapshot(feed_a) if hasattr(engine, "feed_snapshot") else None
        snap_b = engine.feed_snapshot(feed_b) if hasattr(engine, "feed_snapshot") else None
        if snap_a is None or snap_b is None:
            return None

        pa = snap_a.price if snap_a.price > 0 else (snap_a.bid + snap_a.ask) / 2
        pb = snap_b.price if snap_b.price > 0 else (snap_b.bid + snap_b.ask) / 2
        if pa <= 0 or pb <= 0:
            return None

        prices_a.append(pa)
        prices_b.append(pb)
        ticks_since_recalib += 1

        min_samples = max(20, config.recalibrate_every)
        if len(prices_a) < min_samples:
            return None

        # Recalibrate periodically
        if ticks_since_recalib >= config.recalibrate_every:
            ticks_since_recalib = 0
            try:
                from horizon._horizon import cointegration_test, signal_half_life, spread_zscore

                ratio, resid, adf = cointegration_test(
                    list(prices_a), list(prices_b),
                )
                hl = signal_half_life(resid)

                # Check half-life bounds
                if hl < config.min_half_life or hl > config.max_half_life:
                    logger.debug(
                        "Stat arb %s/%s: half-life %.1f outside bounds [%.1f, %.1f]",
                        market_a, market_b, hl,
                        config.min_half_life, config.max_half_life,
                    )
                    return None

                hedge_ratio = ratio
                residuals = resid
                last_half_life = hl
                last_adf = adf
            except Exception as e:
                logger.debug("Stat arb recalibration failed: %s", e)
                return None

        if not residuals:
            return None

        # Compute current residual and z-score
        current_resid = pa - hedge_ratio * pb
        residuals_with_current = residuals + [current_resid]

        from horizon._horizon import spread_zscore
        z = spread_zscore(residuals_with_current, config.lookback)

        # Generate signal
        signal = "hold"
        if position_state == "flat":
            if z > config.entry_zscore:
                signal = "long_b_short_a"  # Residual too high, short spread
            elif z < -config.entry_zscore:
                signal = "long_a_short_b"  # Residual too low, long spread
        else:
            if abs(z) < config.exit_zscore:
                signal = "exit"
            elif abs(z) > config.stop_zscore:
                signal = "stop"

        result = StatArbResult(
            pair=config.pair,
            hedge_ratio=hedge_ratio,
            zscore=z,
            half_life=last_half_life,
            adf_stat=last_adf,
            signal=signal,
        )
        ctx.params["last_stat_arb"] = result

        if auto_execute and signal not in ("hold",):
            now = time.monotonic()
            if now - last_exec_time >= cooldown:
                try:
                    _execute_stat_signal(engine, result, size, hedge_ratio, market_a, market_b, pa, pb)
                    last_exec_time = now
                    if signal in ("exit", "stop"):
                        position_state = "flat"
                    else:
                        position_state = signal
                    logger.info(
                        "Stat arb %s/%s: z=%.2f signal=%s hl=%.1f",
                        market_a, market_b, z, signal, last_half_life,
                    )
                except Exception as e:
                    logger.warning("Stat arb execution failed: %s", e)

        return None

    _scanner.__name__ = "stat_arb"
    return _scanner


def _execute_stat_signal(
    engine, result: StatArbResult, size: float, hedge_ratio: float,
    market_a: str, market_b: str, price_a: float, price_b: float,
) -> None:
    """Execute a stat arb signal."""
    from horizon._horizon import OrderSide, Side

    size_b = size * abs(hedge_ratio)
    if result.signal == "long_a_short_b":
        engine.submit_order(market_a, Side.Yes, OrderSide.Buy, size, price_a)
        engine.submit_order(market_b, Side.Yes, OrderSide.Sell, size_b, price_b)
    elif result.signal == "long_b_short_a":
        engine.submit_order(market_b, Side.Yes, OrderSide.Buy, size_b, price_b)
        engine.submit_order(market_a, Side.Yes, OrderSide.Sell, size, price_a)
    elif result.signal in ("exit", "stop"):
        # Close positions by canceling open orders and submitting market exits
        engine.cancel_market(market_a)
        engine.cancel_market(market_b)
        for pos in engine.positions():
            if pos.market_id == market_a and pos.size > 0:
                engine.submit_market_order(market_a, pos.side, OrderSide.Sell, pos.size)
            elif pos.market_id == market_b and pos.size > 0:
                engine.submit_market_order(market_b, pos.side, OrderSide.Sell, pos.size)
