"""Tests for WZRD momentum routing strategy."""

from __future__ import annotations

import asyncio
import sys
import os

# Add parent to path for direct import
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import wzrd_momentum_strategy as wms


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

MOCK_MOMENTUM = {
    "generated_at": "2026-03-15T12:00:00Z",
    "count": 3,
    "contract_version": "wzrd.momentum.v1",
    "models": [
        {
            "model": "Qwen/Qwen3.5-9B",
            "platform": "huggingface",
            "trend": "surging",
            "score": 0.92,
            "confidence": "normal",
            "action": "pre_warm_urgent",
            "capabilities": ["code", "chat", "reasoning", "text"],
        },
        {
            "model": "meta-llama/Llama-3.3-70B-Instruct",
            "platform": "huggingface",
            "trend": "stable",
            "score": 0.45,
            "confidence": "normal",
            "action": "maintain",
            "capabilities": ["code", "chat", "text"],
        },
        {
            "model": "mistralai/mistral-inference",
            "platform": "github",
            "trend": "cooling",
            "score": 0.12,
            "confidence": "normal",
            "action": "consider_deprovision",
            "capabilities": ["code", "chat"],
        },
    ],
}


def make_deployment(model_name: str, litellm_model: str) -> dict:
    return {
        "model_name": model_name,
        "litellm_params": {"model": litellm_model},
        "model_info": {"id": f"id-{model_name}"},
    }


DEPLOYMENTS = [
    make_deployment("qwen-9b", "openrouter/qwen/qwen-3.5-9b"),
    make_deployment("llama-70b", "openrouter/meta-llama/llama-3.3-70b-instruct"),
    make_deployment("mistral", "openrouter/mistralai/mistral-large"),
]


class MockRouter:
    def __init__(self, model_list: list):
        self.model_list = model_list
        self._custom_strategy = None

    def set_custom_routing_strategy(self, strategy):
        self._custom_strategy = strategy
        self.get_available_deployment = strategy.get_available_deployment
        self.async_get_available_deployment = strategy.async_get_available_deployment


# ---------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------

def test_score_surging_beats_stable():
    """Surging model should score higher than stable."""
    surging_score, surging_signal = wms._score_deployment(
        DEPLOYMENTS[0], MOCK_MOMENTUM,
        {"qwen-9b": ["Qwen/Qwen3.5-9B"]}, 0,
    )
    stable_score, stable_signal = wms._score_deployment(
        DEPLOYMENTS[1], MOCK_MOMENTUM,
        {"llama-70b": ["meta-llama/Llama-3.3-70B-Instruct"]}, 1,
    )

    assert surging_signal is not None
    assert stable_signal is not None
    assert surging_score > stable_score, f"{surging_score} should be > {stable_score}"


def test_score_stable_beats_cooling():
    """Stable model should score higher than cooling."""
    stable_score, _ = wms._score_deployment(
        DEPLOYMENTS[1], MOCK_MOMENTUM,
        {"llama-70b": ["meta-llama/Llama-3.3-70B-Instruct"]}, 0,
    )
    cooling_score, _ = wms._score_deployment(
        DEPLOYMENTS[2], MOCK_MOMENTUM,
        {"mistral": ["mistralai/mistral-inference"]}, 1,
    )
    assert stable_score > cooling_score


def test_score_no_signal_returns_order_bias():
    """Deployment with no matching signal gets small negative score based on position."""
    score, signal = wms._score_deployment(
        make_deployment("unknown-model", "some/unknown"), MOCK_MOMENTUM, {}, 3,
    )
    assert signal is None
    assert score < 0  # Order bias


def test_score_empty_momentum():
    """Empty momentum data should give order-based scores."""
    score, signal = wms._score_deployment(DEPLOYMENTS[0], {}, {}, 0)
    assert signal is None
    assert score == 0.0  # index 0 → -0.001 * 0 = 0


def test_confidence_weighting():
    """Low confidence should reduce score impact."""
    low_conf = dict(MOCK_MOMENTUM)
    low_conf["models"] = [{
        **MOCK_MOMENTUM["models"][0],
        "confidence": "low",
    }]

    full_score, _ = wms._score_deployment(
        DEPLOYMENTS[0], MOCK_MOMENTUM,
        {"qwen-9b": ["Qwen/Qwen3.5-9B"]}, 0,
    )
    low_score, _ = wms._score_deployment(
        DEPLOYMENTS[0], low_conf,
        {"qwen-9b": ["Qwen/Qwen3.5-9B"]}, 0,
    )
    assert full_score > low_score, f"Full confidence {full_score} should beat low {low_score}"


def test_insufficient_confidence_zeroes_score():
    """Insufficient confidence should zero out all signal components."""
    insuf = dict(MOCK_MOMENTUM)
    insuf["models"] = [{
        **MOCK_MOMENTUM["models"][0],
        "confidence": "insufficient",
    }]
    score, signal = wms._score_deployment(
        DEPLOYMENTS[0], insuf,
        {"qwen-9b": ["Qwen/Qwen3.5-9B"]}, 0,
    )
    assert signal is not None
    assert score == 0.0


# ---------------------------------------------------------------------------
# Payload contract guard
# ---------------------------------------------------------------------------

def test_payload_contract_accepts_valid_payload():
    """Valid payload should pass contract guard."""
    assert wms._payload_contract_ok(MOCK_MOMENTUM) is True


def test_payload_contract_rejects_missing_version():
    """Missing both signal_version and contract_version should fail contract guard."""
    bad = dict(MOCK_MOMENTUM)
    bad.pop("contract_version", None)
    bad.pop("signal_version", None)
    assert wms._payload_contract_ok(bad) is False


def test_payload_contract_accepts_legacy_signal_version():
    """Legacy signal_version (int) should still pass contract guard."""
    legacy = dict(MOCK_MOMENTUM)
    legacy.pop("contract_version", None)
    legacy["signal_version"] = 2
    assert wms._payload_contract_ok(legacy) is True


def test_payload_contract_rejects_missing_model_fields():
    """Missing required model-level fields should fail contract guard."""
    bad = dict(MOCK_MOMENTUM)
    bad["models"] = [dict(MOCK_MOMENTUM["models"][0])]
    bad["models"][0].pop("trend", None)
    assert wms._payload_contract_ok(bad) is False


# ---------------------------------------------------------------------------
# Model matching
# ---------------------------------------------------------------------------

def test_alias_map_exact_match():
    """Alias map should find exact match."""
    signal = wms._find_signal(
        DEPLOYMENTS[0], MOCK_MOMENTUM["models"],
        {"qwen-9b": ["Qwen/Qwen3.5-9B"]},
    )
    assert signal is not None
    assert signal["model"] == "Qwen/Qwen3.5-9B"


def test_auto_match_from_litellm_params():
    """Should auto-match using litellm_params.model provider path."""
    signal = wms._find_signal(
        DEPLOYMENTS[1], MOCK_MOMENTUM["models"], {},
    )
    assert signal is not None
    assert signal["model"] == "meta-llama/Llama-3.3-70B-Instruct"


def test_no_match_returns_none():
    """Unmatched deployment returns None."""
    signal = wms._find_signal(
        make_deployment("totally-unknown", "foo/bar/baz"),
        MOCK_MOMENTUM["models"], {},
    )
    assert signal is None


def test_extract_model_slugs():
    """Should extract multiple searchable slugs from deployment."""
    slugs = wms._extract_model_slugs(DEPLOYMENTS[0])
    assert "qwen-9b" in slugs
    assert "openrouter/qwen/qwen-3.5-9b" in slugs
    assert "qwen/qwen-3.5-9b" in slugs
    assert "qwen-3.5-9b" in slugs


# ---------------------------------------------------------------------------
# Strategy integration
# ---------------------------------------------------------------------------

def test_strategy_picks_surging():
    """Strategy should pick the surging model over stable and cooling."""
    wms.clear_cache()
    # Monkey-patch fetch to use mock data
    original_fetch = wms._fetch_momentum
    wms._fetch_momentum = lambda url, timeout=5.0: MOCK_MOMENTUM

    try:
        router = MockRouter(DEPLOYMENTS)
        strategy = wms.WZRDMomentumStrategy(
            router,
            alias_map={
                "qwen-9b": ["Qwen/Qwen3.5-9B"],
                "llama-70b": ["meta-llama/Llama-3.3-70B-Instruct"],
                "mistral": ["mistralai/mistral-inference"],
            },
        )
        # Request model group that has all three as candidates (no exact match)
        result = strategy.get_available_deployment("all-models")
        # Should pick qwen-9b (surging) from all deployments
        assert result is not None
        assert result["model_name"] == "qwen-9b"
    finally:
        wms._fetch_momentum = original_fetch


def test_strategy_exact_group_single():
    """Single deployment in model group returns it directly (no API call)."""
    wms.clear_cache()
    router = MockRouter(DEPLOYMENTS)
    strategy = wms.WZRDMomentumStrategy(router)

    result = strategy.get_available_deployment("qwen-9b")
    assert result is not None
    assert result["model_name"] == "qwen-9b"


def test_strategy_empty_model_list():
    """Empty model list returns None."""
    router = MockRouter([])
    strategy = wms.WZRDMomentumStrategy(router)
    result = strategy.get_available_deployment("anything")
    assert result is None


def test_async_strategy():
    """Async path should return same result as sync."""
    wms.clear_cache()
    original_fetch = wms._fetch_momentum
    wms._fetch_momentum = lambda url, timeout=5.0: MOCK_MOMENTUM

    try:
        router = MockRouter(DEPLOYMENTS)
        strategy = wms.WZRDMomentumStrategy(
            router,
            alias_map={
                "qwen-9b": ["Qwen/Qwen3.5-9B"],
                "llama-70b": ["meta-llama/Llama-3.3-70B-Instruct"],
                "mistral": ["mistralai/mistral-inference"],
            },
        )
        result = asyncio.get_event_loop().run_until_complete(
            strategy.async_get_available_deployment("all-models")
        )
        assert result is not None
        assert result["model_name"] == "qwen-9b"
    finally:
        wms._fetch_momentum = original_fetch


def test_graceful_degradation():
    """When WZRD is down, should return first deployment."""
    wms.clear_cache()
    original_fetch = wms._fetch_momentum
    wms._fetch_momentum = lambda url, timeout=5.0: {}

    try:
        router = MockRouter(DEPLOYMENTS)
        strategy = wms.WZRDMomentumStrategy(router)
        result = strategy.get_available_deployment("all-models")
        assert result is not None
        assert result["model_name"] == "qwen-9b"  # First in list
    finally:
        wms._fetch_momentum = original_fetch


# ---------------------------------------------------------------------------
# Register convenience
# ---------------------------------------------------------------------------

def test_register():
    """register() should wire up the strategy."""
    wms.clear_cache()
    router = MockRouter(DEPLOYMENTS)
    strategy = wms.register(router)
    assert isinstance(strategy, wms.WZRDMomentumStrategy)
    assert router._custom_strategy is strategy


# ---------------------------------------------------------------------------
# Cache
# ---------------------------------------------------------------------------

def test_cache_clear():
    """clear_cache() should empty the module cache."""
    wms._cache["test"] = (0.0, {"test": True})
    assert "test" in wms._cache
    wms.clear_cache()
    assert "test" not in wms._cache


# ---------------------------------------------------------------------------
# Run
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    tests = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)]
    passed = 0
    failed = 0
    for test in tests:
        try:
            test()
            print(f"  PASS  {test.__name__}")
            passed += 1
        except Exception as e:
            print(f"  FAIL  {test.__name__}: {e}")
            failed += 1
    print(f"\n{passed} passed, {failed} failed out of {len(tests)} tests")
    sys.exit(1 if failed else 0)
