"""HTTP client for the Ollama local LLM inference server.

Only Ollama is supported out of the box, but the interface is minimal enough
that alternative backends can be added by subclassing :class:`BaseLLMClient`.
"""

from __future__ import annotations

import json
import logging
from abc import ABC, abstractmethod
from typing import Any, cast

import requests

__all__ = ["BaseLLMClient", "OllamaClient"]

logger = logging.getLogger(__name__)


class BaseLLMClient(ABC):
    """Abstract base for LLM clients.

    Implement :meth:`generate` to plug in a different inference backend.
    """

    @abstractmethod
    def generate(self, prompt: str) -> dict[str, Any] | None:
        """Send *prompt* to the model and return the parsed JSON response.

        Args:
            prompt: The full prompt string to send.

        Returns:
            A ``dict`` decoded from the model's JSON response, or ``None``
            if the request or JSON parsing failed.
        """


class OllamaClient(BaseLLMClient):
    """Thin HTTP client for the `Ollama <https://ollama.com>`_ REST API.

    Args:
        base_url: Base URL of the Ollama server
                  (default: ``"http://localhost:11434"``).
        model:    Model tag to use (default: ``"qwen2.5:7b-instruct"``).
        timeout:  Request timeout in seconds (default: ``120``).

    Example::

        client = OllamaClient(model="llama3:8b")
        result = client.generate("Extract entities from: ...")
        # result is a dict or None
    """

    def __init__(
        self,
        base_url: str = "http://localhost:11434",
        model: str = "qwen2.5:7b-instruct",
        timeout: int = 120,
    ) -> None:
        self._url = f"{base_url.rstrip('/')}/api/generate"
        self.model = model
        self.timeout = timeout

    def generate(self, prompt: str) -> dict[str, Any] | None:
        """Call the Ollama ``/api/generate`` endpoint and return parsed JSON.

        Args:
            prompt: The full prompt string.

        Returns:
            Decoded ``dict`` from the model's JSON-formatted response, or
            ``None`` on any network or parsing error.
        """
        payload: dict[str, Any] = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "format": "json",
        }
        try:
            response = requests.post(self._url, json=payload, timeout=self.timeout)
            response.raise_for_status()
            parsed = json.loads(response.json()["response"])
            return cast(dict[str, Any], parsed)
        except requests.RequestException as exc:
            logger.error("LLM request failed: %s", exc)
        except (KeyError, json.JSONDecodeError, ValueError) as exc:
            logger.error("LLM response parsing failed: %s", exc)
        return None
