import asyncio
import json
import urllib.error
import urllib.request
from typing import Any

from .ai import AIClientError


class OllamaError(AIClientError):
    pass


class OllamaClient:
    provider = "local"
    display_name = "Ollama"

    def __init__(self, base_url: str) -> None:
        self.base_url = base_url.rstrip("/")

    def normalize_model(self, model: str) -> str:
        return model.strip()

    async def list_models(self) -> list[str]:
        return await asyncio.to_thread(self._list_models)

    async def chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        json_mode: bool = False,
    ) -> str:
        return await asyncio.to_thread(self._chat, model, messages, json_mode)

    def _list_models(self) -> list[str]:
        data = self._request_json("GET", "/api/tags", None, timeout=15)
        models = data.get("models", [])
        return sorted(model["name"] for model in models if model.get("name"))

    def _chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        json_mode: bool,
    ) -> str:
        payload = {
            "model": model,
            "messages": messages,
            "stream": False,
        }
        if json_mode:
            payload["format"] = "json"

        data = self._request_json("POST", "/api/chat", payload, timeout=180)
        content = data.get("message", {}).get("content")
        if isinstance(content, str) and content.strip():
            return content.strip()

        fallback = data.get("response")
        if isinstance(fallback, str) and fallback.strip():
            return fallback.strip()

        raise OllamaError("Ollama returned an empty response.")

    def _request_json(
        self,
        method: str,
        path: str,
        payload: dict[str, Any] | None,
        timeout: int,
    ) -> dict[str, Any]:
        url = f"{self.base_url}{path}"
        data = json.dumps(payload).encode("utf-8") if payload is not None else None
        request = urllib.request.Request(
            url,
            data=data,
            method=method,
            headers={"Content-Type": "application/json"},
        )

        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                raw = response.read().decode("utf-8")
        except urllib.error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")[:400]
            raise OllamaError(f"Ollama HTTP {exc.code}: {body}") from exc
        except urllib.error.URLError as exc:
            raise OllamaError(f"Ollama is unavailable at {self.base_url}: {exc.reason}") from exc
        except TimeoutError as exc:
            raise OllamaError("Ollama request timed out.") from exc

        try:
            parsed = json.loads(raw)
        except json.JSONDecodeError as exc:
            raise OllamaError("Ollama returned invalid JSON.") from exc

        if not isinstance(parsed, dict):
            raise OllamaError("Ollama returned an unexpected response.")
        return parsed