"""
WZRD Momentum Routing Strategy for LiteLLM.

Adds demand-aware routing to any LiteLLM Router. Models with accelerating
adoption get preferred; models losing traction get deprioritized.

Quick start:
    from litellm import Router
    from wzrd_momentum_strategy import register

    router = Router(model_list=[...])
    register(router)

    # Every subsequent router.completion() call routes via WZRD momentum.
"""

from __future__ import annotations

import re
import sys
import time
import logging
from typing import Any, Dict, List, Optional, Tuple, Union

try:
    from litellm.types.router import CustomRoutingStrategyBase
except ImportError:
    class CustomRoutingStrategyBase:  # type: ignore[no-redef]
        def get_available_deployment(self, *a: Any, **kw: Any) -> Any: ...
        async def async_get_available_deployment(self, *a: Any, **kw: Any) -> Any: ...

logger = logging.getLogger("wzrd_momentum")

WZRD_MOMENTUM_URL = "https://api.twzrd.xyz/v1/signals/momentum"
_CACHE_TTL = 300  # 5 min — momentum updates every 5 min server-side
_cache: Dict[str, Tuple[float, Dict[str, Any]]] = {}

TREND_SCORE: Dict[str, float] = {
    "surging": 3.0,
    "accelerating": 2.0,
    "stable": 0.0,
    "insufficient_history": 0.0,
    "decelerating": -1.0,
    "cooling": -2.0,
}

CONFIDENCE_WEIGHT: Dict[str, float] = {
    "normal": 1.0,
    "low": 0.5,
    "insufficient": 0.0,
}

REQUIRED_MODEL_FIELDS = {
    "model",
    "trend",
    "score",
    "confidence",
}


# ---------------------------------------------------------------------------
# Signal fetching
# ---------------------------------------------------------------------------

def _payload_contract_ok(data: Dict[str, Any]) -> bool:
    """Validate minimum API contract to guard against schema drift."""
    has_version = (
        isinstance(data.get("signal_version"), int)
        or isinstance(data.get("contract_version"), str)
    )
    if not has_version:
        return False
    models = data.get("models")
    if not isinstance(models, list):
        return False
    for model in models:
        if not isinstance(model, dict):
            return False
        if not REQUIRED_MODEL_FIELDS.issubset(model.keys()):
            return False
    return True

def _fetch_momentum(url: str, timeout: float = 5.0) -> Dict[str, Any]:
    """Fetch momentum signals with cache. Never raises."""
    now = time.monotonic()
    cached = _cache.get(url)
    if cached and (now - cached[0]) < _CACHE_TTL:
        return cached[1]

    data: Dict[str, Any] = {}
    try:
        import httpx
        resp = httpx.get(url, timeout=timeout)
        resp.raise_for_status()
        data = resp.json()
    except Exception:
        try:
            import requests as req_lib
            resp = req_lib.get(url, timeout=timeout)  # type: ignore[assignment]
            resp.raise_for_status()
            data = resp.json()
        except Exception as e:
            logger.debug("wzrd: momentum unavailable (%s), using deployment order", e)
            return {}

    if not isinstance(data, dict):
        return {}
    if not _payload_contract_ok(data):
        logger.warning("wzrd: momentum payload contract mismatch, using deployment order fallback")
        return {}
    _cache[url] = (now, data)
    return data


def clear_cache() -> None:
    """Clear the momentum signal cache. Useful for testing."""
    _cache.clear()


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

def _normalize(name: str) -> str:
    """Lowercase, collapse whitespace, strip."""
    return re.sub(r"\s+", " ", name.strip().lower())


def _extract_model_slugs(deployment: Dict[str, Any]) -> List[str]:
    """Extract searchable model name slugs from a LiteLLM deployment dict.

    Tries model_name, litellm_params.model (with provider prefix stripped),
    and the final path segment.
    """
    slugs: List[str] = []

    model_name = deployment.get("model_name", "")
    if model_name:
        slugs.append(_normalize(model_name))

    litellm_model = ""
    params = deployment.get("litellm_params")
    if isinstance(params, dict):
        litellm_model = params.get("model", "")

    if litellm_model:
        slugs.append(_normalize(litellm_model))
        # Strip provider prefix: "openrouter/qwen/qwen-3.5-9b" -> "qwen/qwen-3.5-9b"
        parts = litellm_model.split("/", 1)
        if len(parts) == 2:
            slugs.append(_normalize(parts[1]))
        # Last segment only: "qwen-3.5-9b"
        last = litellm_model.rsplit("/", 1)[-1]
        if last:
            slugs.append(_normalize(last))

    return list(dict.fromkeys(slugs))  # dedupe, preserve order


def _find_signal(
    deployment: Dict[str, Any],
    models: List[Dict[str, Any]],
    alias_map: Dict[str, List[str]],
) -> Optional[Dict[str, Any]]:
    """Find the best matching WZRD signal for a deployment."""
    model_name = deployment.get("model_name", "")

    # 1. Explicit alias map (highest priority)
    wzrd_names = alias_map.get(model_name, [])
    if wzrd_names:
        for wzrd_name in wzrd_names:
            wzrd_lower = _normalize(wzrd_name)
            for m in models:
                if _normalize(m.get("model", "")) == wzrd_lower:
                    return m
        # Substring fallback for alias map entries
        for wzrd_name in wzrd_names:
            wzrd_lower = _normalize(wzrd_name)
            for m in models:
                if wzrd_lower in _normalize(m.get("model", "")):
                    return m

    # 2. Auto-match using deployment slugs
    slugs = _extract_model_slugs(deployment)
    for slug in slugs:
        for m in models:
            m_norm = _normalize(m.get("model", ""))
            if slug == m_norm:
                return m
    # Substring match (less precise, still useful)
    for slug in slugs:
        if len(slug) < 4:  # skip very short slugs to avoid false positives
            continue
        for m in models:
            m_norm = _normalize(m.get("model", ""))
            if slug in m_norm or m_norm in slug:
                return m

    return None


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

def _score_deployment(
    deployment: Dict[str, Any],
    momentum_data: Dict[str, Any],
    alias_map: Dict[str, List[str]],
    index: int,
) -> Tuple[float, Optional[Dict[str, Any]]]:
    """Score a deployment. Returns (score, matched_signal)."""
    models = momentum_data.get("models", [])
    if not models:
        return (-0.001 * index, None)

    signal = _find_signal(deployment, models, alias_map)
    if signal is None:
        return (-0.001 * index, None)

    trend = signal.get("trend", signal.get("velocity_trend", "stable"))
    momentum = signal.get("score", signal.get("momentum_score", 0.0))
    delta_pct = signal.get("velocity_delta_pct", 0.0)
    confidence = signal.get("confidence", signal.get("history_confidence", "insufficient"))

    cw = CONFIDENCE_WEIGHT.get(confidence, 0.0)
    trend_component = TREND_SCORE.get(trend, 0.0) * cw
    momentum_component = min(max(momentum, 0.0), 1.0) * 0.3 * cw
    delta_component = max(-2.0, min(2.0, delta_pct / 100.0)) * 0.25 * cw

    return (trend_component + momentum_component + delta_component, signal)


# ---------------------------------------------------------------------------
# Strategy
# ---------------------------------------------------------------------------

class WZRDMomentumStrategy(CustomRoutingStrategyBase):
    """Routes LiteLLM requests to models with the highest WZRD momentum.

    Fetches attention velocity signals from the WZRD API, scores each
    deployment in the Router's model_list. Models with accelerating
    adoption (surging/accelerating) get preferred. Models losing
    traction (cooling/decelerating) get deprioritized.

    Falls back to first deployment if WZRD is unreachable.
    """

    def __init__(
        self,
        router: Any,
        wzrd_url: str = WZRD_MOMENTUM_URL,
        alias_map: Optional[Dict[str, List[str]]] = None,
        cache_ttl: int = _CACHE_TTL,
    ):
        self.router = router
        self.wzrd_url = wzrd_url
        self.alias_map = alias_map or {}
        global _CACHE_TTL
        _CACHE_TTL = cache_ttl

    def _select_deployment(
        self,
        model: str,
        messages: Optional[List[Dict[str, str]]] = None,
        input: Optional[Union[str, List]] = None,
        specific_deployment: Optional[bool] = False,
        request_kwargs: Optional[Dict] = None,
    ) -> Optional[Dict]:
        model_list = getattr(self.router, "model_list", [])
        if not model_list:
            return None

        candidates = [d for d in model_list if d.get("model_name", "") == model]
        if not candidates:
            candidates = model_list

        if len(candidates) == 1:
            return candidates[0]

        momentum = _fetch_momentum(self.wzrd_url)

        scored = []
        for i, deployment in enumerate(candidates):
            score, signal = _score_deployment(deployment, momentum, self.alias_map, i)
            scored.append((score, i, deployment, signal))

        scored.sort(key=lambda x: (x[0], -x[1]), reverse=True)
        best = scored[0]

        if momentum.get("models"):
            trend = best[3].get("velocity_trend", "?") if best[3] else "no-signal"
            logger.info(
                "wzrd: %s (score=%.3f, trend=%s) from %d candidates",
                best[2].get("model_name", "?"), best[0], trend, len(candidates),
            )

        return best[2]

    def get_available_deployment(
        self,
        model: str,
        messages: Optional[List[Dict[str, str]]] = None,
        input: Optional[Union[str, List]] = None,
        specific_deployment: Optional[bool] = False,
        request_kwargs: Optional[Dict] = None,
    ) -> Optional[Dict]:
        return self._select_deployment(model, messages, input, specific_deployment, request_kwargs)

    async def async_get_available_deployment(
        self,
        model: str,
        messages: Optional[List[Dict[str, str]]] = None,
        input: Optional[Union[str, List]] = None,
        specific_deployment: Optional[bool] = False,
        request_kwargs: Optional[Dict] = None,
    ) -> Optional[Dict]:
        return self._select_deployment(model, messages, input, specific_deployment, request_kwargs)


# ---------------------------------------------------------------------------
# Convenience
# ---------------------------------------------------------------------------

def register(
    router: Any,
    wzrd_url: str = WZRD_MOMENTUM_URL,
    alias_map: Optional[Dict[str, List[str]]] = None,
    cache_ttl: int = 300,
) -> WZRDMomentumStrategy:
    """One-liner registration. Returns the strategy instance for inspection.

    Usage:
        from wzrd_momentum_strategy import register
        strategy = register(router, alias_map={"qwen-9b": ["Qwen/Qwen3.5-9B"]})
    """
    strategy = WZRDMomentumStrategy(router, wzrd_url=wzrd_url, alias_map=alias_map, cache_ttl=cache_ttl)
    router.set_custom_routing_strategy(strategy)
    return strategy
