"""LLM Gateway -- multi-provider abstraction.

Supports Ollama (local, free), Claude (Anthropic API), and OpenRouter.
"""
import time
from dataclasses import dataclass
from typing import Any, Protocol, runtime_checkable

import httpx


@dataclass
class CompletionResult:
    """Result from an LLM completion."""
    content: str
    model: str
    provider: str
    tokens_used: int
    latency_ms: int
    cost_estimate: float
    tool_calls: list[dict] | None = None


@runtime_checkable
class LLMGateway(Protocol):
    """Protocol for LLM providers."""

    @property
    def provider(self) -> str: ...

    @property
    def model_name(self) -> str: ...

    async def check_available(self) -> bool: ...

    async def complete(self, messages: list[dict], **kwargs) -> CompletionResult: ...


class OllamaGateway:
    """Ollama local LLM gateway. Free, always available if Ollama is running."""

    def __init__(self, base_url: str = "http://localhost:11434", model: str = "mistral-nemo"):
        self._base_url = base_url.rstrip("/")
        self._model = model

    @property
    def provider(self) -> str:
        return "ollama"

    @property
    def model_name(self) -> str:
        return self._model

    async def check_available(self) -> bool:
        try:
            async with httpx.AsyncClient(timeout=3.0) as client:
                resp = await client.get(f"{self._base_url}/api/tags")
                return resp.status_code == 200
        except (httpx.ConnectError, httpx.TimeoutException):
            return False

    async def complete(self, messages: list[dict], tools: list[dict] | None = None, **kwargs) -> CompletionResult:
        start = time.monotonic()
        payload = {"model": self._model, "messages": messages, "stream": False, **kwargs}
        if tools:
            payload["tools"] = tools

        async with httpx.AsyncClient(timeout=120.0) as client:
            resp = await client.post(
                f"{self._base_url}/api/chat",
                json=payload,
            )
            resp.raise_for_status()
            data = resp.json()

        elapsed_ms = int((time.monotonic() - start) * 1000)
        message = data.get("message", {})
        content = message.get("content", "")
        tokens = data.get("eval_count", 0) + data.get("prompt_eval_count", 0)

        # Parse tool calls if present
        raw_tool_calls = message.get("tool_calls")
        tool_calls = None
        if raw_tool_calls:
            tool_calls = []
            for tc in raw_tool_calls:
                func = tc.get("function", {})
                tool_calls.append({
                    "function": {
                        "name": func.get("name", ""),
                        "arguments": func.get("arguments", {}),
                    }
                })

        return CompletionResult(
            content=content,
            model=self._model,
            provider="ollama",
            tokens_used=tokens,
            latency_ms=elapsed_ms,
            cost_estimate=0.0,
            tool_calls=tool_calls,
        )

    async def list_models(self) -> list[dict[str, Any]]:
        async with httpx.AsyncClient(timeout=5.0) as client:
            resp = await client.get(f"{self._base_url}/api/tags")
            resp.raise_for_status()
            return resp.json().get("models", [])


class ClaudeGateway:
    """Anthropic Claude API gateway. Requires anthropic package."""

    def __init__(self, api_key: str | None = None, model: str = "claude-sonnet-4-6"):
        self._api_key = api_key
        self._model = model
        self._client = None

    @property
    def provider(self) -> str:
        return "claude"

    @property
    def model_name(self) -> str:
        return self._model

    async def check_available(self) -> bool:
        if not self._api_key:
            return False
        try:
            import anthropic
            return True
        except ImportError:
            return False

    async def complete(self, messages: list[dict], **kwargs) -> CompletionResult:
        import anthropic

        start = time.monotonic()
        client = anthropic.AsyncAnthropic(api_key=self._api_key)

        # Separate system messages from user/assistant messages
        system_parts = [m["content"] for m in messages if m["role"] == "system"]
        chat_messages = [m for m in messages if m["role"] != "system"]
        system_text = "\n\n".join(system_parts) if system_parts else None

        response = await client.messages.create(
            model=self._model,
            max_tokens=kwargs.get("max_tokens", 4096),
            system=system_text or anthropic.NOT_GIVEN,
            messages=chat_messages,
        )

        elapsed_ms = int((time.monotonic() - start) * 1000)
        content = response.content[0].text if response.content else ""
        tokens = (response.usage.input_tokens or 0) + (response.usage.output_tokens or 0)

        return CompletionResult(
            content=content,
            model=self._model,
            provider="claude",
            tokens_used=tokens,
            latency_ms=elapsed_ms,
            cost_estimate=tokens * 0.000015,  # rough estimate
        )
