"""llama.cpp integration — local inference via llama.cpp server or library."""

from __future__ import annotations

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

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

logger = logging.getLogger(__name__)


@dataclass
class LlamaCppConfig:
    """Configuration for llama.cpp client."""

    base_url: str = "http://127.0.0.1:8080"
    timeout: float = 600.0
    model_path: str | None = None
    n_ctx: int = 2048
    n_gpu_layers: int = 0
    n_threads: int = 4
    verify_ssl: bool = True


class LlamaCppClient(ProviderABC):
    """llama.cpp client for local model inference."""

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

    @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:
        return self.http.request(method, path, **kwargs)

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

    def list_models(self) -> list[dict[str, Any]]:
        models: list[dict[str, Any]] = []
        response = self._request("GET", "/api/tokenize")
        if response.status_code == 200:
            try:
                data = response.json()
                if isinstance(data, dict):
                    models.append({"model": data.get("model", "unknown")})
            except Exception:  # noqa: BLE001
                logger.debug("Failed to parse models from llama.cpp")
        return models

    def info(self) -> dict[str, Any]:
        response = self._request("GET", "/")
        if response.status_code == 200:
            return response.json()
        return {}

    def chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        **kwargs: Any,
    ) -> ChatResult:
        temperature = kwargs.get("temperature")
        max_tokens = kwargs.get("max_tokens")
        top_p = kwargs.get("top_p")
        top_k = kwargs.get("top_k")
        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 max_tokens is not None:
            body.setdefault("options", {})
            body["options"]["max_tokens"] = max_tokens
        if top_p is not None:
            body.setdefault("options", {})
            body["options"]["top_p"] = top_p
        if top_k is not None:
            body.setdefault("options", {})
            body["options"]["top_k"] = top_k
        if extra:
            body.update(extra)

        logger.info("llama.cpp chat: model=%s", model)

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

        if stream:
            return self._stream_chat(model, data)
        choices = data.get("choices", [])
        if choices:
            return choices[0].get("message", {}).get("content", "")
        return ""

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

    def completion(
        self,
        model: str,
        prompt: str,
        *,
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | 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 max_tokens is not None:
            body.setdefault("options", {})
            body["options"]["max_tokens"] = max_tokens
        if top_p is not None:
            body.setdefault("options", {})
            body["options"]["top_p"] = top_p
        if extra:
            body.update(extra)

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

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

        if stream:
            return self._stream_completion(model, data)
        choices = data.get("choices", [])
        if choices:
            return choices[0].get("text", "")
        return ""

    def _stream_completion(
        self, model: str, data: Any
    ) -> Iterator[dict[str, Any]]:
        for chunk in data:
            yield {
                "model": chunk.get("model", model),
                "text": chunk.get("text", ""),
                "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("llama.cpp embed: model=%s", model)

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

    def tokenize(self, model: str, text: str) -> list[int]:
        response = self._request("POST", "/api/tokenize", json={
            "model": model,
            "content": text,
        })
        response.raise_for_status()
        data = response.json()
        tokens = data.get("tokens", [])
        return [int(t) for t in tokens]

    def detokenize(self, model: str, tokens: list[int]) -> str:
        response = self._request("POST", "/api/detokenize", json={
            "model": model,
            "tokens": tokens,
        })
        response.raise_for_status()
        data = response.json()
        return data.get("text", "")

    def _load_library(self) -> Any:
        if self._lib_loaded:
            try:
                from llama_cpp import Llama  # type: ignore[import-not-found]
                return Llama
            except ImportError:
                return None
        try:
            from llama_cpp import Llama  # type: ignore[import-not-found]
            self._lib_loaded = True
            return Llama
        except ImportError:
            return None

    def load_library_model(
        self,
        model_path: str | None = None,
        n_ctx: int | None = None,
        n_gpu_layers: int | None = None,
        n_threads: int | None = None,
    ) -> Any:
        Llama_cls = self._load_library()
        if Llama_cls is None:
            raise ImportError(
                "llama-cpp-python is required for library mode. "
                "Install with: pip install llama-cpp-python"
            )

        kwargs: dict[str, Any] = {
            "model_path": model_path or self.config.model_path or "",
            "n_ctx": n_ctx or self.config.n_ctx,
            "n_gpu_layers": n_gpu_layers or self.config.n_gpu_layers,
            "n_threads": n_threads or self.config.n_threads,
            "verbose": False,
        }
        kwargs = {k: v for k, v in kwargs.items() if v}

        logger.info("Loading llama.cpp model: %s", kwargs["model_path"])
        return Llama_cls(**kwargs)

    async def achat(
        self,
        model: str,
        messages: list[dict[str, str]],
        *,
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | 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 max_tokens is not None:
            body.setdefault("options", {})
            body["options"]["max_tokens"] = max_tokens
        if top_p is not None:
            body.setdefault("options", {})
            body["options"]["top_p"] = top_p
        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),
                "content": chunk.get("content", ""),
                "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()

        choices = data.get("choices", [])
        if choices:
            return choices[0].get("message", {}).get("content", "")
        return ""

    async def acompletion(
        self,
        model: str,
        prompt: str,
        *,
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | 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 max_tokens is not None:
            body.setdefault("options", {})
            body["options"]["max_tokens"] = max_tokens
        if top_p is not None:
            body.setdefault("options", {})
            body["options"]["top_p"] = top_p
        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/completion", json=body)
        response.raise_for_status()
        data = response.json()

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

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

        choices = data.get("choices", [])
        if choices:
            return choices[0].get("text", "")
        return ""

    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/embeddings", 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:
            info = self.info()
            return {
                "healthy": True,
                "provider": "llama.cpp",
                "base_url": self.config.base_url,
                "server_info": info,
            }
        except Exception as exc:  # noqa: BLE001
            return {"healthy": False, "provider": "llama.cpp", "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():
                    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