from __future__ import annotations

"""Ollama provider backed by its OpenAI-compatible endpoint."""

from dataclasses import dataclass
from typing import AsyncIterator

from agentrspy.core import (
    CompletionRequest,
    CompletionResponse,
    LlmProvider,
    StreamChunk,
)

from .openai import OpenAiBuilder, OpenAiProvider


@dataclass(slots=True)
class OllamaProvider(LlmProvider):
    """Ollama provider backed by the OpenAI-compatible API."""

    inner: OpenAiProvider

    @classmethod
    def builder(cls) -> "OllamaBuilder":
        return OllamaBuilder()

    @classmethod
    def new(cls, model: str) -> "OllamaProvider":
        return cls.builder().model(model).build()

    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        return await self.inner.complete(request)

    async def stream(self, request: CompletionRequest) -> AsyncIterator[StreamChunk]:
        return await self.inner.stream(request)

    def name(self) -> str:
        return "ollama"


@dataclass(slots=True)
class OllamaBuilder:
    """Builder for `OllamaProvider`."""

    model_name: str = "llama3.2"
    base_url_value: str = "http://localhost:11434/v1"

    def model(self, model: str) -> "OllamaBuilder":
        self.model_name = model
        return self

    def base_url(self, base_url: str) -> "OllamaBuilder":
        self.base_url_value = base_url
        return self

    def build(self) -> OllamaProvider:
        inner = (
            OpenAiBuilder()
            .api_key("ollama")
            .base_url(self.base_url_value)
            .model(self.model_name)
            .build()
        )
        return OllamaProvider(inner=inner)
