"""Ollama integration — local model serving via Ollama API."""

from __future__ import annotations

import json
import logging
from collections.abc import AsyncIterator, Iterator
from dataclasses import dataclass
from typing import Any

from cyberbrain26.model_layer.provider_interface import ChatResult, ProviderABC, ProviderProtocol
import httpx

logger = logging.getLogger(__name__)


@dataclass
class OllamaConfig:
    """Configuration for Ollama client."""

    base_url: str = "http://127.0.0.1:11434"
    timeout: float = 60.0
    verify_ssl: bool = True


class OllamaClient(ProviderABC):
    """Ollama client for local model inference."""

    def __init__(self, config: OllamaConfig | None = None) -> None:
        self.config = config or OllamaConfig()
        self._http: httpx.Client | None = None
        self._async_http: httpx.AsyncClient | None = None

    @property
    def http(self) -> httpx.Client:
        if self._http is None:
            self._http = httpx.Client(
                base_url=self.config.base_url,
                timeout=self.config.timeout,
                verify=self.config.verify_ssl,
            )
        return self._http

    @property
    def async_http(self) -> httpx.AsyncClient:
        if self._async_http is None:
            self._async_http = httpx.AsyncClient(
                base_url=self.config.base_url,
                timeout=self.config.timeout,
                verify=self.config.verify_ssl,
            )
        return self._async_http

    def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response:
        kwargs.setdefault("timeout", self.config.timeout)
        return self.http.request(method, path, **kwargs)

    async def _arequest(self, method: str, path: str, **kwargs: Any) -> httpx.Response:
        kwargs.setdefault("timeout", self.config.timeout)
        return await self.async_http.request(method, path, **kwargs)

    def list_models(self) -> list[dict[str, Any]]:
        response = self._request("GET", "/api/tags")
        response.raise_for_status()
        data = response.json()
        return data.get("models", [])

    def list_running(self) -> list[dict[str, Any]]:
        response = self._request("GET", "/api/ps")
        response.raise_for_status()
        data = response.json()
        return data.get("models", [])

    def show_model(self, model_name: str) -> dict[str, Any]:
        response = self._request("POST", "/api/show", json={"model": model_name})
        response.raise_for_status()
        return response.json()

    def pull_model(
        self, model_name: str, *, stream: bool = False
    ) -> dict[str, Any] | Iterator[dict[str, Any]]:
        response = self._request(
            "POST", "/api/pull", json={"model": model_name}, stream=stream
        )
        if stream:
            return self._pull_stream(model_name, response)
        response.raise_for_status()
        return response.json()

    def _pull_stream(
        self, model_name: str, response: httpx.Response
    ) -> Iterator[dict[str, Any]]:
        for line in response.iter_lines():
            if line:
                yield json.loads(line)

    def chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        **kwargs: Any,
    ) -> ChatResult:
        temperature = kwargs.get("temperature")
        num_predict = kwargs.get("num_predict")
        stream = kwargs.get("stream", False)
        extra = kwargs.get("extra")
        body: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": stream,
        }
        if temperature is not None:
            body.setdefault("options", {})
            body["options"]["temperature"] = temperature
        if num_predict is not None:
            body.setdefault("options", {})
            body["options"]["num_predict"] = num_predict
        if extra:
            body.update(extra)

        logger.info("Ollama chat: model=%s, messages=%d", model, len(messages))

        response = self._request("POST", "/api/chat", json=body)
        response.raise_for_status()
        data = response.json()

        if stream:
            return self._stream_chat(model, data)
        return data.get("message", {}).get("content", "")

    def _stream_chat(
        self, model: str, data: Any
    ) -> Iterator[dict[str, Any]]:
        for chunk in data:
            yield {
                "model": chunk.get("model", model),
                "message": chunk.get("message", {}),
                "done": chunk.get("done", False),
            }

    def completion(
        self,
        model: str,
        prompt: str,
        *,
        temperature: float | None = None,
        num_predict: int | None = None,
        stream: bool = False,
        extra: dict[str, Any] | None = None,
    ) -> str | Iterator[dict[str, Any]]:
        body: dict[str, Any] = {
            "model": model,
            "prompt": prompt,
            "stream": stream,
        }
        if temperature is not None:
            body.setdefault("options", {})
            body["options"]["temperature"] = temperature
        if num_predict is not None:
            body.setdefault("options", {})
            body["options"]["num_predict"] = num_predict
        if extra:
            body.update(extra)

        logger.info("Ollama completion: model=%s", model)

        response = self._request("POST", "/api/generate", json=body)
        response.raise_for_status()
        data = response.json()

        if stream:
            return self._stream_completion(model, data)
        return data.get("response", "")

    def _stream_completion(
        self, model: str, data: Any
    ) -> Iterator[dict[str, Any]]:
        for chunk in data:
            yield {
                "model": chunk.get("model", model),
                "response": chunk.get("response", ""),
                "done": chunk.get("done", False),
            }

    def embed(
        self,
        model: str,
        texts: list[str],
        **kwargs: Any,
    ) -> list[list[float]]:
        extra = kwargs.get("extra")
        body: dict[str, Any] = {
            "model": model,
            "input": texts,
        }
        if extra:
            body.update(extra)

        logger.info("Ollama embed: model=%s", model)

        response = self._request("POST", "/api/embed", json=body)
        response.raise_for_status()
        data = response.json()
        return data.get("embeddings", [])

    async def achat(
        self,
        model: str,
        messages: list[dict[str, str]],
        *,
        temperature: float | None = None,
        num_predict: int | None = None,
        stream: bool = False,
        extra: dict[str, Any] | None = None,
    ) -> str | AsyncIterator[dict[str, Any]]:
        body: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": stream,
        }
        if temperature is not None:
            body.setdefault("options", {})
            body["options"]["temperature"] = temperature
        if num_predict is not None:
            body.setdefault("options", {})
            body["options"]["num_predict"] = num_predict
        if extra:
            body.update(extra)

        if stream:
            return self._astream_chat(model, body)
        return await self._achat_nonstream(model, body)

    async def _astream_chat(
        self,
        model: str,
        body: dict[str, Any],
    ) -> AsyncIterator[dict[str, Any]]:
        response = await self._arequest("POST", "/api/chat", json=body)
        response.raise_for_status()
        data = response.json()

        for chunk in data:
            yield {
                "model": chunk.get("model", model),
                "message": chunk.get("message", {}),
                "done": chunk.get("done", False),
            }

    async def _achat_nonstream(
        self,
        model: str,
        body: dict[str, Any],
    ) -> str:
        response = await self._arequest("POST", "/api/chat", json=body)
        response.raise_for_status()
        data = response.json()
        return data.get("message", {}).get("content", "")

    async def acompletion(
        self,
        model: str,
        prompt: str,
        *,
        temperature: float | None = None,
        num_predict: int | None = None,
        stream: bool = False,
        extra: dict[str, Any] | None = None,
    ) -> str | AsyncIterator[dict[str, Any]]:
        body: dict[str, Any] = {
            "model": model,
            "prompt": prompt,
            "stream": stream,
        }
        if temperature is not None:
            body.setdefault("options", {})
            body["options"]["temperature"] = temperature
        if num_predict is not None:
            body.setdefault("options", {})
            body["options"]["num_predict"] = num_predict
        if extra:
            body.update(extra)

        if stream:
            return self._astream_completion(model, body)
        return await self._acompletion_nonstream(model, body)

    async def _astream_completion(
        self,
        model: str,
        body: dict[str, Any],
    ) -> AsyncIterator[dict[str, Any]]:
        response = await self._arequest("POST", "/api/generate", json=body)
        response.raise_for_status()
        data = response.json()

        for chunk in data:
            yield {
                "model": chunk.get("model", model),
                "response": chunk.get("response", ""),
                "done": chunk.get("done", False),
            }

    async def _acompletion_nonstream(
        self,
        model: str,
        body: dict[str, Any],
    ) -> str:
        response = await self._arequest("POST", "/api/generate", json=body)
        response.raise_for_status()
        data = response.json()
        return data.get("response", "")

    async def aembed(
        self,
        model: str,
        texts: list[str],
        *,
        extra: dict[str, Any] | None = None,
    ) -> list[list[float]]:
        body: dict[str, Any] = {
            "model": model,
            "input": texts,
        }
        if extra:
            body.update(extra)

        response = await self._arequest("POST", "/api/embed", json=body)
        response.raise_for_status()
        data = response.json()
        return data.get("embeddings", [])

    def is_available(self) -> bool:
        try:
            response = self._request("GET", "/")
            return response.status_code == 200
        except Exception:  # noqa: BLE001
            return False

    def health_check(self) -> dict[str, Any]:
        try:
            models = self.list_models()
            running = self.list_running()
            return {
                "healthy": True,
                "provider": "ollama",
                "base_url": self.config.base_url,
                "models_available": len(models),
                "models_running": len(running),
            }
        except Exception as exc:  # noqa: BLE001
            return {"healthy": False, "provider": "ollama", "error": str(exc)}

    def close(self) -> None:
        if self._http is not None:
            self._http.close()
        if self._async_http is not None:
            try:
                import asyncio

                loop = asyncio.get_event_loop()
                if loop.is_running():
                    # Schedule closing; can't await in sync context
                    loop.create_task(self._async_http.aclose())
                else:
                    loop.run_until_complete(self._async_http.aclose())
            except Exception:  # noqa: BLE001, S110
                pass
            finally:
                self._async_http = None