"""OpenAI-compatible and Anthropic chat models as FGAO LLMs (sync and native async).

``OpenAILLM`` also serves any OpenAI-compatible endpoint (vLLM, TGI, Ollama,
LM Studio, ...) through ``base_url``. Clients are created lazily, so a process
that only uses the async API never opens a sync client and vice versa.
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional

from ..rag.prompts import SYSTEM_PROMPT


def _messages(system: Optional[str], prompt: str) -> List[Dict[str, str]]:
    return ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": prompt}]


class OpenAILLM:
    def __init__(self, model: str, *, client: Any = None, aclient: Any = None, base_url: Optional[str] = None,
                 api_key: Optional[str] = None, max_tokens: int = 128, temperature: float = 0.0,
                 system: Optional[str] = SYSTEM_PROMPT, timeout: Optional[float] = None,
                 max_retries: Optional[int] = None) -> None:
        self.model, self.max_tokens, self.temperature, self.system = model, max_tokens, temperature, system
        self._kw = {k: v for k, v in (("base_url", base_url), ("api_key", api_key), ("timeout", timeout),
                                      ("max_retries", max_retries)) if v is not None}
        self._client, self._aclient = client, aclient

    @property
    def client(self):
        if self._client is None:
            from openai import OpenAI

            self._client = OpenAI(**self._kw)
        return self._client

    @property
    def aclient(self):
        if self._aclient is None:
            from openai import AsyncOpenAI

            self._aclient = AsyncOpenAI(**self._kw)
        return self._aclient

    def _args(self, prompt: str) -> Dict[str, Any]:
        return {"model": self.model, "messages": _messages(self.system, prompt), "max_tokens": self.max_tokens,
                "temperature": self.temperature}

    def complete(self, prompt: str) -> str:
        out = self.client.chat.completions.create(**self._args(prompt))
        return (out.choices[0].message.content or "").strip()

    async def acomplete(self, prompt: str) -> str:
        out = await self.aclient.chat.completions.create(**self._args(prompt))
        return (out.choices[0].message.content or "").strip()


class AnthropicLLM:
    def __init__(self, model: str, *, client: Any = None, aclient: Any = None, api_key: Optional[str] = None,
                 max_tokens: int = 128, temperature: float = 0.0, system: Optional[str] = SYSTEM_PROMPT,
                 timeout: Optional[float] = None, max_retries: Optional[int] = None) -> None:
        self.model, self.max_tokens, self.temperature, self.system = model, max_tokens, temperature, system
        self._kw = {k: v for k, v in (("api_key", api_key), ("timeout", timeout), ("max_retries", max_retries))
                    if v is not None}
        self._client, self._aclient = client, aclient

    @property
    def client(self):
        if self._client is None:
            from anthropic import Anthropic

            self._client = Anthropic(**self._kw)
        return self._client

    @property
    def aclient(self):
        if self._aclient is None:
            from anthropic import AsyncAnthropic

            self._aclient = AsyncAnthropic(**self._kw)
        return self._aclient

    def _args(self, prompt: str) -> Dict[str, Any]:
        kw: Dict[str, Any] = {"model": self.model, "max_tokens": self.max_tokens, "temperature": self.temperature,
                              "messages": [{"role": "user", "content": prompt}]}
        if self.system:
            kw["system"] = self.system
        return kw

    def complete(self, prompt: str) -> str:
        out = self.client.messages.create(**self._args(prompt))
        return "".join(getattr(b, "text", "") for b in out.content).strip()

    async def acomplete(self, prompt: str) -> str:
        out = await self.aclient.messages.create(**self._args(prompt))
        return "".join(getattr(b, "text", "") for b in out.content).strip()


__all__ = ["OpenAILLM", "AnthropicLLM"]