"""
Ollama client for OomLlama.
"""

import httpx
from typing import Optional, Dict, Any, List, Generator


class OllamaClient:
    """
    HTTP client for Ollama API.

    Handles connection to local or remote Ollama instances.
    """

    def __init__(
        self,
        base_url: str = "http://localhost:11434",
        timeout: float = 60.0
    ):
        """
        Initialize Ollama client.

        Args:
            base_url: Ollama API URL
            timeout: Request timeout in seconds
        """
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    def generate(
        self,
        model: str,
        prompt: str,
        system: Optional[str] = None,
        temperature: float = 0.7,
        top_p: float = 0.9,
        max_tokens: int = 512,
        stream: bool = False
    ) -> Dict[str, Any]:
        """
        Generate text with Ollama.

        Args:
            model: Model name (e.g., "qwen2.5:7b")
            prompt: User prompt
            system: System prompt
            temperature: Sampling temperature
            top_p: Top-p sampling
            max_tokens: Maximum tokens to generate
            stream: Stream response

        Returns:
            Response dict with 'response' key
        """
        payload = {
            "model": model,
            "prompt": prompt,
            "stream": stream,
            "options": {
                "temperature": temperature,
                "top_p": top_p,
                "num_predict": max_tokens
            }
        }

        if system:
            payload["system"] = system

        with httpx.Client(timeout=self.timeout) as client:
            response = client.post(
                f"{self.base_url}/api/generate",
                json=payload
            )
            response.raise_for_status()
            return response.json()

    def chat(
        self,
        model: str,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 512
    ) -> Dict[str, Any]:
        """
        Chat with Ollama.

        Args:
            model: Model name
            messages: List of {"role": "user/assistant", "content": "..."}
            temperature: Sampling temperature
            max_tokens: Maximum tokens

        Returns:
            Response dict
        """
        payload = {
            "model": model,
            "messages": messages,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens
            }
        }

        with httpx.Client(timeout=self.timeout) as client:
            response = client.post(
                f"{self.base_url}/api/chat",
                json=payload
            )
            response.raise_for_status()
            return response.json()

    def list_models(self) -> List[str]:
        """List available models."""
        with httpx.Client(timeout=10.0) as client:
            response = client.get(f"{self.base_url}/api/tags")
            response.raise_for_status()
            data = response.json()
            return [m["name"] for m in data.get("models", [])]

    def is_available(self) -> bool:
        """Check if Ollama is reachable."""
        try:
            with httpx.Client(timeout=3.0) as client:
                response = client.get(f"{self.base_url}/api/tags")
                return response.status_code == 200
        except:
            return False

    def pull(self, model: str) -> bool:
        """
        Pull a model.

        Args:
            model: Model to pull (e.g., "qwen2.5:7b")

        Returns:
            True if successful
        """
        with httpx.Client(timeout=300.0) as client:
            response = client.post(
                f"{self.base_url}/api/pull",
                json={"name": model, "stream": False}
            )
            return response.status_code == 200
