"""LLM client supporting Ollama with a plain HTTP fallback.

The :class:`OllamaLLM` class first tries the official ``ollama`` Python
package; if it is not installed it falls back to a direct HTTP call to the
Ollama REST API using ``requests`` (which is already a required dependency).
"""

from __future__ import annotations

import logging
import time
from typing import Optional, Tuple

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Result type
# ---------------------------------------------------------------------------


class LLMResult:
    """Response returned by :class:`OllamaLLM`."""

    def __init__(
        self,
        text: str,
        generate_time: float,
        eval_count: Optional[int] = None,
    ) -> None:
        self.text = text
        self.generate_time = generate_time  # seconds
        self.eval_count = eval_count        # tokens generated (from Ollama)
        self.words = len(text.split())

    def __repr__(self) -> str:
        return (
            f"LLMResult(words={self.words}, "
            f"generate_time={self.generate_time:.1f}s, "
            f"eval_count={self.eval_count})"
        )


# ---------------------------------------------------------------------------
# Ollama client
# ---------------------------------------------------------------------------


class OllamaLLM:
    """Thin wrapper around the Ollama generate API.

    Parameters
    ----------
    model:
        Ollama model tag (e.g. ``"llama3.2"``, ``"mistral"``,
        ``"gpt-oss:20b"``).
    host:
        Base URL of the Ollama server (default: ``http://localhost:11434``).
    num_ctx:
        Context-window size in tokens passed to Ollama (default: 131 072).

    Examples
    --------
    >>> from neuromind.alzheimers.llm import OllamaLLM
    >>> llm = OllamaLLM(model="llama3.2")
    >>> result = llm.generate("What is Alzheimer's disease?", max_tokens=200)
    >>> print(result.text)
    """

    def __init__(
        self,
        model: str = "llama3.2",
        host: str = "http://localhost:11434",
        num_ctx: int = 131_072,
    ) -> None:
        self.model = model
        self.host = host.rstrip("/")
        self.num_ctx = num_ctx
        self._has_package = self._check_ollama_package()

    # ------------------------------------------------------------------
    # Package detection
    # ------------------------------------------------------------------

    @staticmethod
    def _check_ollama_package() -> bool:
        try:
            import ollama  # noqa: F401
            return True
        except ImportError:
            return False

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def generate(
        self,
        prompt: str,
        temperature: float = 0.3,
        max_tokens: int = 2000,
    ) -> LLMResult:
        """Generate a completion for *prompt*.

        Parameters
        ----------
        prompt:
            Full prompt string (system instructions + context + question).
        temperature:
            Sampling temperature.  0 = near-deterministic, 1 = creative.
        max_tokens:
            Maximum number of tokens to generate (``num_predict`` in Ollama).

        Returns
        -------
        :class:`LLMResult`
        """
        t0 = time.time()
        if self._has_package:
            text, eval_count = self._via_package(prompt, temperature, max_tokens)
        else:
            text, eval_count = self._via_http(prompt, temperature, max_tokens)
        elapsed = time.time() - t0

        return LLMResult(text=text, generate_time=elapsed, eval_count=eval_count)

    # ------------------------------------------------------------------
    # Backends
    # ------------------------------------------------------------------

    def _via_package(
        self, prompt: str, temperature: float, max_tokens: int
    ) -> Tuple[str, Optional[int]]:
        import ollama

        try:
            response = ollama.generate(
                model=self.model,
                prompt=prompt,
                options={
                    "temperature": temperature,
                    "num_predict": max_tokens,
                    "num_ctx": self.num_ctx,
                },
            )
        except Exception as exc:
            msg = str(exc).lower()
            if any(k in msg for k in ("connect", "refused", "connection")):
                raise ConnectionError(
                    f"Cannot connect to Ollama at {self.host}.\n"
                    "Make sure Ollama is running: run 'ollama serve' in a terminal,\n"
                    f"then verify the model is available: 'ollama pull {self.model}'"
                ) from exc
            raise
        # The ollama package returns a dict-like object
        text = (
            response.get("response", "") if isinstance(response, dict)
            else getattr(response, "response", "")
        )
        eval_count = (
            response.get("eval_count") if isinstance(response, dict)
            else getattr(response, "eval_count", None)
        )
        return text, eval_count

    def _via_http(
        self, prompt: str, temperature: float, max_tokens: int
    ) -> Tuple[str, Optional[int]]:
        try:
            import requests
        except ImportError as exc:
            raise ImportError(
                "Install either the 'ollama' package or 'requests' to use the LLM.\n"
                "  pip install ollama\n"
                "  pip install requests"
            ) from exc

        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens,
                "num_ctx": self.num_ctx,
            },
        }
        try:
            resp = requests.post(
                f"{self.host}/api/generate",
                json=payload,
                timeout=300,
            )
        except requests.exceptions.ConnectionError as exc:
            raise ConnectionError(
                f"Cannot connect to Ollama at {self.host}.\n"
                "Make sure Ollama is running: run 'ollama serve' in a terminal,\n"
                f"then verify the model is available: 'ollama pull {self.model}'"
            ) from exc
        resp.raise_for_status()
        data = resp.json()
        return data.get("response", ""), data.get("eval_count")
