from __future__ import annotations

"""Ollama LLM client implementation.

This module provides a local Ollama client that integrates with the neuraparse
LLM abstraction layer. Ollama allows running LLMs locally without API keys.
"""

from dataclasses import dataclass, field
from typing import Any

try:
    import ollama  # type: ignore[import-not-found]
except ImportError:  # pragma: no cover
    ollama = None  # type: ignore[assignment]


@dataclass
class OllamaLLMClient:
    """Ollama local LLM client.

    Uses the Ollama Python SDK to generate completions from locally running models.
    No API key required; assumes Ollama server is running on localhost:11434.

    Args:
        model: Ollama model name (e.g., "llama3.2", "mistral", "phi3").
        base_url: Ollama server URL (default: "http://localhost:11434").
        temperature: Sampling temperature (0.0 to 2.0).
        max_tokens: Maximum tokens in the completion (Ollama uses 'num_predict').
    """

    model: str = "llama3.2"
    base_url: str = "http://localhost:11434"
    temperature: float = 0.7
    max_tokens: int = 1024
    _client: Any = field(init=False, repr=False, default=None)

    def __post_init__(self) -> None:
        """Initialize Ollama client after dataclass initialization."""
        if ollama is None:  # pragma: no cover
            raise RuntimeError(
                "Ollama provider requires 'ollama' package. "
                "Install it with: pip install ollama"
            )

        self._client = ollama.Client(host=self.base_url)

    def generate(self, prompt: str, **kwargs: Any) -> str:
        """Generate a completion for the given prompt.

        Args:
            prompt: Input prompt text.
            **kwargs: Optional overrides for temperature, max_tokens, etc.

        Returns:
            Generated completion text.
        """
        temperature = kwargs.get("temperature", self.temperature)
        max_tokens = kwargs.get("max_tokens", self.max_tokens)

        response = self._client.generate(
            model=self.model,
            prompt=prompt,
            options={
                "temperature": temperature,
                "num_predict": max_tokens,
            },
        )

        return response.get("response", "")

