"""LangChain tools for WZRD velocity signals and oracle inference.

WzrdModelPicker  - free, zero-config model selection via /v1/signals/momentum
WzrdInference    - paid inference through WZRD's oracle via /v1/chat/completions
"""

from __future__ import annotations

import logging
import os
import threading
import time
from typing import Any, Dict, List, Optional, Tuple, Type

import httpx
from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

logger = logging.getLogger("langchain_wzrd")

MOMENTUM_URL = "https://api.twzrd.xyz/v1/signals/momentum"
COMPLETIONS_URL = "https://api.twzrd.xyz/v1/chat/completions"
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(capability: Optional[str] = None) -> Dict[str, Any]:
    """Fetch velocity signal with 60s cache. Never raises."""
    url = MOMENTUM_URL
    if capability and capability != "auto":
        url += f"?capability={capability}"
    now = time.monotonic()
    with _cache_lock:
        cached = _cache.get(url)
        if cached and (now - cached[0]) < CACHE_TTL:
            return cached[1]
    try:
        r = httpx.get(url, timeout=5.0)
        r.raise_for_status()
        data = r.json()
    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 _rank_models(signal: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Sort models 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]


class ModelPickerInput(BaseModel):
    """Input for WzrdModelPicker."""
    task: str = Field(
        default="auto",
        description="Task type: 'code', 'chat', 'reasoning', or 'auto' for best overall.",
    )


class WzrdModelPicker(BaseTool):
    """Pick the best AI model right now based on live velocity signals.

    Returns the top model name, its velocity score, and adoption trend.
    Uses WZRD's free /v1/signals/momentum endpoint. No API key required.
    Use this when the user asks which model to use, wants the fastest-growing
    model, or needs a recommendation for a specific task type.
    """

    name: str = "wzrd_model_picker"
    description: str = (
        "Pick the best AI model based on real-time adoption velocity across 100+ models. "
        "Input a task type (code/chat/reasoning/auto). Returns the top model, "
        "its velocity score, and trend direction. Free, no API key needed."
    )
    args_schema: Type[BaseModel] = ModelPickerInput

    def _run(self, task: str = "auto") -> str:
        capability = task if task in ("code", "chat", "reasoning") else None
        signal = _fetch_signal(capability)
        if not signal:
            return "Unable to fetch velocity signals. The WZRD API may be temporarily unavailable."
        ranked = _rank_models(signal)
        if not ranked:
            return "No models available in the current velocity signal."

        top = ranked[0]
        model = top.get("model", "unknown")
        score = top.get("score", 0.0)
        trend = top.get("trend", "unknown")
        confidence = top.get("confidence", "unknown")

        lines = [
            f"Top model: {model}",
            f"Velocity score: {score:.3f}",
            f"Trend: {trend}",
            f"Confidence: {confidence}",
        ]
        if len(ranked) > 1:
            runners = [m.get("model", "?") for m in ranked[1:4]]
            lines.append(f"Runners-up: {', '.join(runners)}")

        ts = signal.get("generated_at", "")
        if ts:
            lines.append(f"Signal timestamp: {ts}")

        return "\n".join(lines)


class InferenceInput(BaseModel):
    """Input for WzrdInference."""
    prompt: str = Field(description="The prompt to send for completion.")
    model: str = Field(
        default="",
        description="Model to use. Leave empty to auto-select the top velocity model.",
    )


class WzrdInference(BaseTool):
    """Run inference through WZRD's oracle and get a completion.

    Calls /v1/chat/completions with the given prompt. Requires a Bearer token
    or AgentCash auth (set WZRD_API_KEY env var). If no model is specified,
    auto-selects the top velocity model. Use this when the user wants to
    actually generate text using the best available model.
    """

    name: str = "wzrd_inference"
    description: str = (
        "Run inference through WZRD's oracle. Input a prompt and optional model name. "
        "Returns the completion text and routing metadata. Requires WZRD_API_KEY env var. "
        "If no model is given, auto-selects the top velocity model."
    )
    args_schema: Type[BaseModel] = InferenceInput
    api_key: str = ""

    def __init__(self, api_key: str = "", **kwargs: Any):
        super().__init__(api_key=api_key, **kwargs)
        if not self.api_key:
            self.api_key = os.environ.get("WZRD_API_KEY", "")

    def _run(self, prompt: str, model: str = "") -> str:
        if not self.api_key:
            return "WZRD_API_KEY not set. Set it as an env var or pass api_key= to WzrdInference()."

        if not model:
            signal = _fetch_signal()
            ranked = _rank_models(signal)
            model = ranked[0].get("model", "") if ranked else ""
        if not model:
            return "No model available. Provide a model name or check the WZRD signal API."

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
        }

        try:
            r = httpx.post(COMPLETIONS_URL, json=payload, headers=headers, timeout=30.0)
            r.raise_for_status()
            data = r.json()
        except httpx.HTTPStatusError as e:
            return f"Inference failed ({e.response.status_code}): {e.response.text[:200]}"
        except Exception as e:
            return f"Inference request failed: {e}"

        choices = data.get("choices", [])
        text = choices[0].get("message", {}).get("content", "") if choices else ""
        usage = data.get("usage", {})

        lines = [text]
        lines.append(f"\n---\nModel: {data.get('model', model)}")
        if usage:
            lines.append(f"Tokens: {usage.get('prompt_tokens', 0)} in / {usage.get('completion_tokens', 0)} out")

        return "\n".join(lines)
