from collections.abc import AsyncIterator

import httpx

from app.config import settings


class OllamaClient:
    def __init__(self, base_url: str | None = None):
        self._base_url = base_url or settings.ollama_base_url

    async def generate(
        self,
        model: str,
        prompt: str,
        system: str | None = None,
        options: dict | None = None,
    ) -> str:
        payload: dict = {"model": model, "prompt": prompt, "stream": False}
        if system is not None:
            payload["system"] = system
        if options:
            payload["options"] = options

        async with httpx.AsyncClient(base_url=self._base_url, timeout=120.0) as http_client:
            response = await http_client.post("/api/generate", json=payload)
            response.raise_for_status()
            return response.json()["response"]

    async def generate_stream(
        self,
        model: str,
        prompt: str,
        system: str | None = None,
        options: dict | None = None,
    ) -> AsyncIterator[str]:
        """Yields response tokens as Ollama produces them."""
        payload: dict = {"model": model, "prompt": prompt, "stream": True}
        if system is not None:
            payload["system"] = system
        if options:
            payload["options"] = options

        async with httpx.AsyncClient(base_url=self._base_url, timeout=300.0) as http_client:
            async with http_client.stream(
                "POST", "/api/generate", json=payload
            ) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if not line:
                        continue
                    import json

                    chunk = json.loads(line)
                    token = chunk.get("response", "")
                    if token:
                        yield token
                    if chunk.get("done"):
                        return