from __future__ import annotations

import asyncio
import json
from typing import Any

import httpx
import tiktoken

from nexus.protocols.providers.base import AbstractProvider, LLMRequest, LLMResponse, ToolCall
from nexus.resilience.errors import ProviderError


class LocalProvider(AbstractProvider):
    """OpenAI-compatible local provider for Ollama, vLLM, LM Studio, llama-server, and similar APIs."""

    def __init__(
        self,
        model: str,
        base_url: str = "http://localhost:11434/v1",
        api_key: str = "not-required",
        context_window_size: int = 32768,
        timeout: float = 120.0,
        backend: str = "ollama",
    ) -> None:
        self._model = model
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._context_window_size = context_window_size
        self._backend = backend.lower()
        if self._backend not in {"ollama", "vllm", "lmstudio", "llama-server", "openai_compat"}:
            raise ValueError("backend must be one of: ollama, vllm, lmstudio, llama-server, openai_compat")
        self._http = httpx.AsyncClient(base_url=self._base_url, timeout=timeout)
        self._enc = tiktoken.get_encoding("cl100k_base")

    async def complete(self, request: LLMRequest) -> LLMResponse:
        """Call local OpenAI-compatible /chat/completions endpoint."""
        body = {
            "model": self._model,
            "messages": self._to_openai_messages(request),
            "max_tokens": request.max_tokens,
            "temperature": request.temperature,
        }
        if request.tools:
            body["tools"] = [
                {
                    "type": "function",
                    "function": {
                        "name": t.name,
                        "description": t.description,
                        "parameters": t.parameters,
                    },
                }
                for t in request.tools
            ]
        if request.stop_sequences:
            body["stop"] = request.stop_sequences

        headers = {"Authorization": f"Bearer {self._api_key}"}
        retries = 2
        for attempt in range(retries + 1):
            try:
                resp = await self._post_chat(body, headers)
                resp.raise_for_status()
                data = resp.json()
                break
            except httpx.ConnectError as exc:
                if attempt >= retries:
                    raise ProviderError(
                        (
                            f"Local provider connection failed at {self._base_url}. "
                            f"Check that backend '{self._backend}' is running and reachable."
                        ),
                        provider_name=self.provider_name,
                    ) from exc
                await asyncio.sleep(0.5 * (attempt + 1))
            except httpx.HTTPStatusError as exc:
                raise ProviderError(
                    f"Local provider HTTP error: {exc.response.status_code}",
                    provider_name=self.provider_name,
                    status_code=exc.response.status_code,
                ) from exc
            except (json.JSONDecodeError, ValueError) as exc:
                raise ProviderError(f"Invalid JSON response from local provider: {exc}", provider_name=self.provider_name) from exc

        choice = (data.get("choices") or [{}])[0]
        message = choice.get("message") or {}
        usage = data.get("usage") or {}

        tool_calls: list[ToolCall] = []
        for call in message.get("tool_calls") or []:
            fn = call.get("function") or {}
            args_raw = fn.get("arguments") or "{}"
            try:
                args = json.loads(args_raw) if isinstance(args_raw, str) else dict(args_raw)
            except (ValueError, TypeError):
                args = {"_raw": str(args_raw)}
            tool_calls.append(
                ToolCall(
                    tool_call_id=call.get("id", ""),
                    tool_name=fn.get("name", ""),
                    arguments=args,
                )
            )

        content = message.get("content") or ""
        in_tok = int(usage.get("prompt_tokens", 0) or 0)
        out_tok = int(usage.get("completion_tokens", 0) or 0)
        if in_tok == 0 and out_tok == 0:
            in_tok = self.count_tokens("\n".join(m.get("content", "") for m in body["messages"] if isinstance(m.get("content"), str)))
            out_tok = self.count_tokens(content)

        return LLMResponse(
            content=content,
            tool_calls=tool_calls,
            input_tokens=in_tok,
            output_tokens=out_tok,
            model=str(data.get("model", self._model)),
            finish_reason=str(choice.get("finish_reason", "stop")),
            raw=data,
        )

    def _to_openai_messages(self, request: LLMRequest) -> list[dict[str, Any]]:
        messages: list[dict[str, Any]] = []
        if request.system:
            messages.append({"role": "system", "content": request.system})
        for m in request.messages:
            item: dict[str, Any] = {"role": m.role, "content": m.content}
            if m.name:
                item["name"] = m.name
            if m.tool_call_id:
                item["tool_call_id"] = m.tool_call_id
            messages.append(item)
        return messages

    def count_tokens(self, text: str) -> int:
        """Estimate token count using cl100k_base encoder."""
        return len(self._enc.encode(text))

    async def list_models(self) -> list[str]:
        """Return available model names from /models endpoint."""
        try:
            resp = await self._get_models_endpoint()
            if resp.status_code != 200:
                return []
            data = resp.json()
            models = data.get("data") or data.get("models") or []
            out: list[str] = []
            for model in models:
                if isinstance(model, dict):
                    name = model.get("id") or model.get("name")
                    if name:
                        out.append(str(name))
                elif isinstance(model, str):
                    out.append(model)
            return out
        except Exception:
            return []

    async def health_check(self) -> bool:
        """Check provider health via /models without raising."""
        try:
            resp = await self._get_models_endpoint()
            return resp.status_code == 200
        except Exception:
            return False

    async def _post_chat(self, body: dict[str, Any], headers: dict[str, str]) -> httpx.Response:
        """Try common OpenAI-compatible chat completion routes."""
        last_exc: Exception | None = None
        for path in ("/chat/completions", "/v1/chat/completions"):
            try:
                return await self._http.post(path, json=body, headers=headers)
            except httpx.HTTPError as exc:
                last_exc = exc
                continue
        assert last_exc is not None
        raise last_exc

    async def _get_models_endpoint(self) -> httpx.Response:
        """Try common OpenAI-compatible model-list routes."""
        last_exc: Exception | None = None
        for path in ("/models", "/v1/models"):
            try:
                return await self._http.get(path)
            except httpx.HTTPError as exc:
                last_exc = exc
                continue
        assert last_exc is not None
        raise last_exc

    @property
    def context_window(self) -> int:
        return self._context_window_size

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

    @property
    def provider_name(self) -> str:
        return f"local_{self._backend}"
