"""Ollama backend implementation."""

import ollama as ollama_client
from typing import Optional

from .base import LLMBackend


class OllamaBackend(LLMBackend):
    """Ollama API backend for local LLM inference."""

    def __init__(
        self,
        host: str = "localhost:11434",
        model: str = "codellama",
        temperature: float = 0.1,
        max_tokens: int = 500,
    ):
        """Initialize the Ollama backend.

        Args:
            host: Ollama server host.
            model: Model name to use.
            temperature: Generation temperature.
            max_tokens: Maximum tokens to generate.
        """
        self.host = host
        self._model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self._client = None

    @property
    def client(self):
        """Get or create Ollama client."""
        if self._client is None:
            self._client = ollama_client.Client(host=self.host)
        return self._client

    def generate(self, prompt: str) -> str:
        """Generate response using Ollama API.

        Args:
            prompt: The prompt to send.

        Returns:
            Generated response text.
        """
        try:
            response = self.client.generate(
                model=self._model,
                prompt=prompt,
                options={
                    "temperature": self.temperature,
                    "num_predict": self.max_tokens,
                },
            )

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

        except Exception as e:
            raise ConnectionError(f"Ollama API error: {e}")

    def is_available(self) -> bool:
        """Check if Ollama is running and model is available.

        Returns:
            True if backend is available.
        """
        try:
            self.client.ps()
            return True
        except Exception:
            return False

    def get_model_name(self) -> str:
        """Get the current model name.

        Returns:
            Model name string.
        """
        return self._model

    def set_model(self, model: str) -> None:
        """Set the model to use.

        Args:
            model: Name of the model.
        """
        self._model = model

    def close(self) -> None:
        """Clean up client resources."""
        self._client = None