"""
WZRD Earn Callback for LiteLLM.

Reports every successful LiteLLM completion to WZRD so the operator earns CCM.
Pair with wzrd_momentum_strategy for the full loop:
    route via velocity → complete → report pick → earn CCM → claim.

Quick start:
    from litellm import Router
    from wzrd_momentum_strategy import register
    from wzrd_earn_callback import WZRDEarnCallback

    router = Router(model_list=[...])
    register(router)                          # route via momentum
    litellm.success_callback = [WZRDEarnCallback()]  # earn from picks

Requires: pip install wzrd-client  (has Ed25519 auth + claim built in)
"""

from __future__ import annotations

import os
import time
import logging
import threading
from typing import Any, Dict, Optional

logger = logging.getLogger("wzrd_earn")

# Lazy-loaded WZRDAgent — avoids import cost if callback is registered but
# never fires (e.g. during testing).
_agent = None
_agent_lock = threading.Lock()


def _get_agent():
    """Lazy-init the WZRDAgent. Thread-safe, one-time setup."""
    global _agent
    if _agent is not None:
        return _agent
    with _agent_lock:
        if _agent is not None:
            return _agent
        try:
            from wzrd import WZRDAgent
            agent = WZRDAgent.from_env()
            session = agent.authenticate()
            _agent = agent
            logger.info("wzrd_earn: authenticated agent %s", session.pubkey[:12])
        except Exception as e:
            logger.warning("wzrd_earn: agent init failed (%s) — will retry next call", e)
            _agent = None  # None = retry later; False would be permanent
    return _agent


def _extract_model_id(kwargs: Dict[str, Any]) -> Optional[str]:
    """Extract the model identifier from LiteLLM kwargs."""
    # litellm_params.model is the canonical source (includes provider prefix)
    litellm_params = kwargs.get("litellm_params", {})
    model = litellm_params.get("model", "")
    if model:
        return model

    # Fallback: top-level model key
    return kwargs.get("model", None)


def _extract_latency_ms(kwargs: Dict[str, Any], response: Any) -> Optional[int]:
    """Extract response latency from LiteLLM response metadata."""
    # LiteLLM stores timing in response._response_ms or in kwargs
    if hasattr(response, "_response_ms"):
        return int(response._response_ms)

    # Some versions put it in additional_args
    start = kwargs.get("start_time")
    end = kwargs.get("end_time")
    if start and end:
        try:
            delta = end - start
            return int(delta.total_seconds() * 1000)
        except Exception:
            pass

    return None


def _extract_cost(kwargs: Dict[str, Any], response: Any) -> Optional[float]:
    """Extract cost from LiteLLM response (if available)."""
    # response_cost is set by litellm.completion_cost()
    cost = kwargs.get("response_cost")
    if cost is not None:
        try:
            return float(cost)
        except (TypeError, ValueError):
            pass
    return None


class WZRDEarnCallback:
    """LiteLLM CustomLogger that reports completions to WZRD.

    Every successful completion is reported with model_id, latency, and cost.
    The operator earns CCM proportional to report volume and pick quality.

    Env vars:
        WZRD_AGENT_KEYPAIR_PATH — path to Ed25519 keypair JSON (required)
        WZRD_API_BASE_URL       — override API base (default: https://api.twzrd.xyz)
        WZRD_REPORT_TASK        — task_type label (default: "inference")
        WZRD_EARN_ENABLED       — set to "false" to disable reporting (default: "true")
    """

    def __init__(self, task_type: str = "inference"):
        self.task_type = os.environ.get("WZRD_REPORT_TASK", task_type)
        self.enabled = os.environ.get("WZRD_EARN_ENABLED", "true").lower() != "false"
        self._report_count = 0
        self._error_count = 0
        self._last_error: Optional[str] = None

    def log_success_event(self, kwargs: Dict[str, Any], response_obj: Any, start_time: Any, end_time: Any) -> None:
        """Called by LiteLLM after every successful completion."""
        if not self.enabled:
            return

        agent = _get_agent()
        if not agent or agent is False:
            return

        model_id = _extract_model_id(kwargs)
        if not model_id:
            return

        # Pass start_time/end_time from callback args into kwargs for extraction
        if start_time is not None and "start_time" not in kwargs:
            kwargs["start_time"] = start_time
        if end_time is not None and "end_time" not in kwargs:
            kwargs["end_time"] = end_time
        latency_ms = _extract_latency_ms(kwargs, response_obj)
        cost_usd = _extract_cost(kwargs, response_obj)

        # Extract provider from model path (e.g. "openrouter/qwen/..." → "openrouter")
        provider = model_id.split("/")[0] if "/" in model_id else "unknown"

        metadata = {
            "wzrd": {
                "source": "litellm-plugin",
                "policy": "litellm-router",
                "policy_version": "1.0",
                "signal_model": model_id,
                "selected_model": model_id,
                "platform": provider,
            },
            "inference": {
                "provider": provider,
                "provider_model": model_id,
                "graded": False,
            },
        }

        try:
            agent.report(
                model_id=model_id,
                task_type=self.task_type,
                latency_ms=latency_ms,
                cost_usd=cost_usd,
                outcome="success",
                metadata=metadata,
            )
            self._report_count += 1
            if self._report_count % 100 == 0:
                logger.info("wzrd_earn: %d reports sent", self._report_count)
        except Exception as e:
            self._error_count += 1
            self._last_error = str(e)
            if self._error_count <= 3:
                logger.warning("wzrd_earn: report failed (%s)", e)

    def log_failure_event(self, kwargs: Dict[str, Any], response_obj: Any, start_time: Any, end_time: Any) -> None:
        """Called by LiteLLM after a failed completion. Reports as failure."""
        if not self.enabled:
            return

        agent = _get_agent()
        if not agent or agent is False:
            return

        model_id = _extract_model_id(kwargs)
        if not model_id:
            return

        try:
            agent.report(
                model_id=model_id,
                task_type=self.task_type,
                outcome="failure",
            )
        except Exception:
            pass  # failures are best-effort

    @property
    def stats(self) -> Dict[str, Any]:
        """Return reporting stats for health checks."""
        return {
            "reports_sent": self._report_count,
            "errors": self._error_count,
            "last_error": self._last_error,
            "enabled": self.enabled,
        }
