"""WZRD Velocity Callback for LiteLLM.

Drop-in callback: fetches velocity signals, ranks models, reports outcomes.
Zero config required. Set WZRD_PREMIUM=true for enriched signals, provide
WZRD_AGENT_KEYPAIR_PATH to earn CCM from inference reports.

    import litellm
    from litellm_wzrd_momentum import WzrdCallback
    litellm.callbacks = [WzrdCallback()]
"""
from __future__ import annotations
import os, time, logging, threading
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger("wzrd")

MOMENTUM_URL = "https://api.twzrd.xyz/v1/signals/momentum"
PREMIUM_URL = "https://api.twzrd.xyz/v1/signals/momentum/premium"
CACHE_TTL = 60

TREND_BOOST = {"surging": 3.0, "accelerating": 2.0, "stable": 0.0,
               "insufficient_history": 0.0, "decelerating": -1.0, "cooling": -2.0}
CONFIDENCE_WEIGHT = {"normal": 1.0, "low": 0.5, "insufficient": 0.0}
_cache_lock = threading.Lock()
_cache: Dict[str, Tuple[float, Dict[str, Any]]] = {}


def _fetch_signal(premium: bool = False, capability: Optional[str] = None) -> Dict[str, Any]:
    """Fetch velocity signal with 60s cache. Never raises."""
    url = PREMIUM_URL if premium else MOMENTUM_URL
    if capability:
        url += f"?capability={capability}"
    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
        r = httpx.get(url, timeout=5.0); r.raise_for_status(); data = r.json()
    except ImportError:
        try:
            import requests as rl
            r = rl.get(url, timeout=5.0); r.raise_for_status(); data = r.json()  # type: ignore[assignment]
        except Exception as e:
            logger.debug("wzrd: signal fetch failed (%s)", e); return {}
    except Exception as e:
        logger.debug("wzrd: signal fetch failed (%s)", e); return {}
    if not isinstance(data, dict) or not isinstance(data.get("models"), list):
        return {}
    with _cache_lock:
        _cache[url] = (now, data)
    return data


def clear_cache() -> None:
    with _cache_lock:
        _cache.clear()


def _rank_models(signal: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Return signal models sorted by composite velocity score (descending)."""
    scored: List[Tuple[float, int, Dict[str, Any]]] = []
    for i, m in enumerate(signal.get("models", [])):
        cw = CONFIDENCE_WEIGHT.get(m.get("confidence", "insufficient"), 0.0)
        s = TREND_BOOST.get(m.get("trend", "stable"), 0.0) * cw
        s += min(max(m.get("score", 0.0), 0.0), 1.0) * 0.3 * cw
        scored.append((s, i, m))
    scored.sort(key=lambda x: (x[0], -x[1]), reverse=True)
    return [x[2] for x in scored]


def pick_model(signal: Dict[str, Any]) -> Optional[str]:
    """Pick the top-ranked model from a velocity signal response."""
    ranked = _rank_models(signal)
    return ranked[0]["model"] if ranked else None


# -- Lazy agent init (thread-safe, one attempt) ---
_agent: Any = None
_agent_lock = threading.Lock()
_agent_tried = False


def _get_agent() -> Any:
    global _agent, _agent_tried
    if _agent is not None:
        return _agent
    if _agent_tried:
        return None
    with _agent_lock:
        if _agent is not None:
            return _agent
        _agent_tried = True
        try:
            from wzrd import WZRDAgent
            a = WZRDAgent.from_env(); a.authenticate(); _agent = a
            logger.info("wzrd: agent authenticated")
        except Exception as e:
            logger.debug("wzrd: agent init skipped (%s)", e)
    return _agent


def _extract_model(kwargs: Dict[str, Any]) -> Optional[str]:
    lp = kwargs.get("litellm_params", {})
    return (lp.get("model") if isinstance(lp, dict) else None) or kwargs.get("model")


def _extract_latency(response: Any, start_time: Any, end_time: Any) -> Optional[int]:
    if hasattr(response, "_response_ms"):
        return int(response._response_ms)
    if start_time and end_time:
        try:
            return int((end_time - start_time).total_seconds() * 1000)
        except Exception:
            pass
    return None


class WzrdCallback:
    """LiteLLM callback integrating WZRD velocity signals.

    Env vars:
        WZRD_PREMIUM            "true" for enriched signal endpoint
        WZRD_CAPABILITY         filter by capability (code, chat, reasoning, vision)
        WZRD_AGENT_KEYPAIR_PATH Ed25519 keypair path for CCM earning
        WZRD_EARN_ENABLED       "false" to disable reporting
    """

    def __init__(self, premium: Optional[bool] = None, capability: Optional[str] = None):
        self.premium = premium if premium is not None else (
            os.environ.get("WZRD_PREMIUM", "").lower() == "true")
        self.capability = capability or os.environ.get("WZRD_CAPABILITY")
        self.earn_enabled = os.environ.get("WZRD_EARN_ENABLED", "true").lower() != "false"
        self._picks = 0
        self._reports = 0
        self._failures = 0

    def log_pre_api_call(self, model: str, messages: Any, kwargs: Dict[str, Any]) -> None:
        signal = _fetch_signal(premium=self.premium, capability=self.capability)
        if not signal:
            return
        top = pick_model(signal)
        if not top:
            return
        self._picks += 1
        kwargs.setdefault("wzrd_signal", {})
        kwargs["wzrd_signal"]["top_model"] = top
        kwargs["wzrd_signal"]["ranked"] = [m["model"] for m in _rank_models(signal)[:5]]
        if self._picks <= 3 or self._picks % 50 == 0:
            logger.info("wzrd: top=%s for model=%s", top, model)

    def log_success_event(self, kwargs: Dict[str, Any], response_obj: Any,
                          start_time: Any, end_time: Any) -> None:
        if not self.earn_enabled:
            return
        agent = _get_agent()
        if agent is None:
            return
        model_id = _extract_model(kwargs)
        if not model_id:
            return
        provider = model_id.split("/")[0] if "/" in model_id else "unknown"
        try:
            agent.report(
                model_id=model_id, task_type="inference",
                latency_ms=_extract_latency(response_obj, start_time, end_time),
                cost_usd=float(c) if (c := kwargs.get("response_cost")) is not None else None,
                outcome="success",
                metadata={"wzrd": {
                    "source": "litellm-callback", "policy_version": "2.0",
                    "signal_model": kwargs.get("wzrd_signal", {}).get("top_model"),
                    "selected_model": model_id, "platform": provider,
                }},
            )
            self._reports += 1
        except Exception as e:
            if self._reports == 0:
                logger.warning("wzrd: report failed (%s)", e)

    def log_failure_event(self, kwargs: Dict[str, Any], response_obj: Any,
                          start_time: Any, end_time: Any) -> None:
        self._failures += 1
        model_id = _extract_model(kwargs)
        if model_id and self._failures <= 5:
            logger.debug("wzrd: failure model=%s total=%d", model_id, self._failures)
        if not self.earn_enabled:
            return
        agent = _get_agent()
        if agent is None or not model_id:
            return
        try:
            agent.report(model_id=model_id, task_type="inference", outcome="failure")
        except Exception:
            pass

    @property
    def stats(self) -> Dict[str, Any]:
        return {"picks": self._picks, "reports": self._reports, "failures": self._failures,
                "premium": self.premium, "earn_enabled": self.earn_enabled}
