from __future__ import annotations

from datetime import UTC, datetime, timedelta

from arbitrade.strategy.stat_arb import StatArbExperiment, StatArbExperimentConfig


def test_stat_arb_experiment_warmup_then_entry_and_exit() -> None:
    started_at = datetime(2026, 6, 2, 12, 0, tzinfo=UTC)
    experiment = StatArbExperiment(
        StatArbExperimentConfig(
            pair_a="BTC/USD",
            pair_b="ETH/USD",
            lookback_window=5,
            entry_zscore=1.5,
            exit_zscore=0.2,
            max_holding_seconds=0.5,
        )
    )

    # Warmup with nearly stationary spread around 0.
    for idx in range(5):
        signal = experiment.observe(
            price_a=100.0 + (0.02 * idx),
            price_b=100.0,
            observed_at=started_at + timedelta(seconds=idx),
        )

    assert signal.action in {"warmup", "hold"}

    # Large positive spread should trigger short-spread entry.
    entry = experiment.observe(
        price_a=104.0,
        price_b=100.0,
        observed_at=started_at + timedelta(seconds=10),
    )
    assert entry.action == "enter_short_spread"
    assert entry.position == "short"
    assert entry.zscore is not None

    # Mean reversion toward center should trigger exit.
    exit_signal = experiment.observe(
        price_a=100.05,
        price_b=100.0,
        observed_at=started_at + timedelta(seconds=11),
    )
    assert exit_signal.action == "exit_position"
    assert exit_signal.position == "flat"


def test_stat_arb_experiment_rejects_invalid_prices() -> None:
    experiment = StatArbExperiment(
        StatArbExperimentConfig(
            pair_a="BTC/USD",
            pair_b="ETH/USD",
            lookback_window=5,
        )
    )

    at = datetime(2026, 6, 2, 12, 0, tzinfo=UTC)
    try:
        experiment.observe(price_a=0.0, price_b=100.0, observed_at=at)
    except ValueError as exc:
        assert "prices must be > 0" in str(exc)
    else:
        raise AssertionError("Expected ValueError for non-positive price")