"""LLM client -- real model invocation for Hive agent reasoning.

Calls Ollama, OpenAI, Anthropic, or OpenRouter via httpx.
Every agent uses this to think, analyze, and generate.
"""

from __future__ import annotations

import logging
import os
import time
from dataclasses import dataclass, field
from typing import Any

import httpx

logger = logging.getLogger("hive.llm")

# Pricing per 1M tokens (USD), approximate as of Feb 2026
_PRICING: dict[str, dict[str, float]] = {
    "gpt-4o": {"input": 2.50, "output": 10.00},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "claude-sonnet-4-20250514": {"input": 3.00, "output": 15.00},
    "claude-3-5-haiku-20241022": {"input": 0.80, "output": 4.00},
}


@dataclass
class LLMResponse:
    """Response from an LLM call."""

    content: str
    model: str
    provider: str
    input_tokens: int = 0
    output_tokens: int = 0
    latency_ms: int = 0
    cost_usd: float = 0.0
    raw: dict[str, Any] = field(default_factory=dict)


class HiveLLMClient:
    """Unified LLM client for all Hive agents.

    Resolves provider from model name or explicit config.
    Falls back gracefully: Ollama (local) → OpenRouter → OpenAI → Anthropic.
    """

    def __init__(
        self,
        ollama_url: str = "http://localhost:11434",
        timeout: float = 120.0,
    ) -> None:
        self._ollama_url = ollama_url
        self._timeout = timeout
        self._usage_log: list[dict[str, Any]] = []

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def complete(
        self,
        messages: list[dict[str, str]],
        model: str | None = None,
        system: str | None = None,
        temperature: float = 0.3,
        max_tokens: int = 4096,
    ) -> LLMResponse:
        """Send a completion request to the best available provider.

        Resolution order:
        1. If model starts with a known prefix, route to that provider.
        2. Otherwise try Ollama → OpenRouter → OpenAI → Anthropic in order.
        """
        if system:
            messages = [{"role": "system", "content": system}, *messages]

        provider, resolved_model = self._resolve_provider(model)
        start = time.monotonic()

        if provider == "ollama":
            result = await self._call_ollama(messages, resolved_model, temperature, max_tokens)
        elif provider == "openai":
            result = await self._call_openai(messages, resolved_model, temperature, max_tokens)
        elif provider == "anthropic":
            result = await self._call_anthropic(messages, resolved_model, temperature, max_tokens)
        elif provider == "openrouter":
            result = await self._call_openrouter(messages, resolved_model, temperature, max_tokens)
        else:
            msg = f"Unknown provider: {provider}"
            raise ValueError(msg)

        result.latency_ms = int((time.monotonic() - start) * 1000)
        result.cost_usd = self._estimate_cost(resolved_model, result.input_tokens, result.output_tokens)

        self._usage_log.append(
            {
                "model": result.model,
                "provider": result.provider,
                "input_tokens": result.input_tokens,
                "output_tokens": result.output_tokens,
                "latency_ms": result.latency_ms,
                "cost_usd": result.cost_usd,
            }
        )

        return result

    async def is_available(self, provider: str | None = None) -> bool:
        """Check if any (or specific) provider is reachable."""
        if provider == "ollama" or provider is None:
            try:
                async with httpx.AsyncClient(timeout=3.0) as client:
                    r = await client.get(f"{self._ollama_url}/api/tags")
                    if r.status_code == 200:
                        return True
            except Exception:
                if provider == "ollama":
                    return False

        if provider in ("openai", None) and os.getenv("OPENAI_API_KEY"):
            return True
        if provider in ("anthropic", None) and os.getenv("ANTHROPIC_API_KEY"):
            return True
        return bool(provider in ("openrouter", None) and os.getenv("OPENROUTER_API_KEY"))

    def get_usage_summary(self) -> dict[str, Any]:
        """Return aggregated usage stats."""
        total_cost = sum(e["cost_usd"] for e in self._usage_log)
        total_calls = len(self._usage_log)
        total_tokens = sum(e["input_tokens"] + e["output_tokens"] for e in self._usage_log)
        return {
            "total_calls": total_calls,
            "total_tokens": total_tokens,
            "total_cost_usd": round(total_cost, 6),
            "by_provider": self._group_by("provider"),
        }

    # ------------------------------------------------------------------
    # Provider Implementations
    # ------------------------------------------------------------------

    async def _call_ollama(
        self,
        messages: list[dict[str, str]],
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> LLMResponse:
        """Call Ollama's /api/chat endpoint."""
        payload = {
            "model": model,
            "messages": messages,
            "stream": False,
            "options": {"temperature": temperature, "num_predict": max_tokens},
        }
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            r = await client.post(f"{self._ollama_url}/api/chat", json=payload)
            r.raise_for_status()
            data = r.json()

        content = data.get("message", {}).get("content", "")
        return LLMResponse(
            content=content,
            model=model,
            provider="ollama",
            input_tokens=data.get("prompt_eval_count", 0),
            output_tokens=data.get("eval_count", 0),
            raw=data,
        )

    async def _call_openai(
        self,
        messages: list[dict[str, str]],
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> LLMResponse:
        """Call OpenAI-compatible /v1/chat/completions."""
        api_key = os.environ["OPENAI_API_KEY"]
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            r = await client.post("https://api.openai.com/v1/chat/completions", json=payload, headers=headers)
            r.raise_for_status()
            data = r.json()

        choice = data["choices"][0]["message"]
        usage = data.get("usage", {})
        return LLMResponse(
            content=choice.get("content", ""),
            model=model,
            provider="openai",
            input_tokens=usage.get("prompt_tokens", 0),
            output_tokens=usage.get("completion_tokens", 0),
            raw=data,
        )

    async def _call_anthropic(
        self,
        messages: list[dict[str, str]],
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> LLMResponse:
        """Call Anthropic /v1/messages endpoint."""
        api_key = os.environ["ANTHROPIC_API_KEY"]

        # Anthropic requires system in a separate field
        system_msg = ""
        chat_messages = []
        for m in messages:
            if m["role"] == "system":
                system_msg += m["content"] + "\n"
            else:
                chat_messages.append(m)

        payload: dict[str, Any] = {
            "model": model,
            "messages": chat_messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if system_msg.strip():
            payload["system"] = system_msg.strip()

        headers = {
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
            "Content-Type": "application/json",
        }
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            r = await client.post("https://api.anthropic.com/v1/messages", json=payload, headers=headers)
            r.raise_for_status()
            data = r.json()

        content_blocks = data.get("content", [])
        content = "".join(b.get("text", "") for b in content_blocks if b.get("type") == "text")
        usage = data.get("usage", {})
        return LLMResponse(
            content=content,
            model=model,
            provider="anthropic",
            input_tokens=usage.get("input_tokens", 0),
            output_tokens=usage.get("output_tokens", 0),
            raw=data,
        )

    async def _call_openrouter(
        self,
        messages: list[dict[str, str]],
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> LLMResponse:
        """Call OpenRouter (OpenAI-compatible)."""
        api_key = os.environ["OPENROUTER_API_KEY"]
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            r = await client.post("https://openrouter.ai/api/v1/chat/completions", json=payload, headers=headers)
            r.raise_for_status()
            data = r.json()

        choice = data["choices"][0]["message"]
        usage = data.get("usage", {})
        return LLMResponse(
            content=choice.get("content", ""),
            model=model,
            provider="openrouter",
            input_tokens=usage.get("prompt_tokens", 0),
            output_tokens=usage.get("completion_tokens", 0),
            raw=data,
        )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _resolve_provider(self, model: str | None) -> tuple[str, str]:
        """Resolve provider and model name from input.

        Returns (provider, model_name).
        """
        if model is None:
            model = self._default_model()

        # Explicit provider prefix: "ollama:llama3", "openai:gpt-4o"
        if ":" in model:
            provider, name = model.split(":", 1)
            return provider, name

        # Known model prefixes
        if model.startswith("gpt-"):
            return "openai", model
        if model.startswith("claude-"):
            return "anthropic", model
        if "/" in model:
            # OpenRouter format: "meta-llama/llama-3-8b"
            return "openrouter", model

        # Default to Ollama for unrecognized models (likely local)
        return "ollama", model

    def _default_model(self) -> str:
        """Pick the best default model from available providers."""
        if os.getenv("OPENAI_API_KEY"):
            return "gpt-4o-mini"
        if os.getenv("ANTHROPIC_API_KEY"):
            return "claude-3-5-haiku-20241022"
        if os.getenv("OPENROUTER_API_KEY"):
            return "meta-llama/llama-3-8b-instruct"
        # Ollama fallback -- assume a common model
        return "ollama:llama3"

    def _estimate_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """Estimate cost in USD."""
        pricing = _PRICING.get(model)
        if not pricing:
            return 0.0
        return (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000

    def _group_by(self, key: str) -> dict[str, Any]:
        """Group usage log by a key."""
        groups: dict[str, dict[str, Any]] = {}
        for entry in self._usage_log:
            k = entry[key]
            if k not in groups:
                groups[k] = {"calls": 0, "tokens": 0, "cost_usd": 0.0}
            groups[k]["calls"] += 1
            groups[k]["tokens"] += entry["input_tokens"] + entry["output_tokens"]
            groups[k]["cost_usd"] += entry["cost_usd"]
        return groups
