"""OpenAI / Ollama LLM provider. Ollama exposes an OpenAI-compatible API."""

from __future__ import annotations

from collections.abc import AsyncIterator

try:
    import openai
except ModuleNotFoundError:
    openai = None  # type: ignore[assignment]

from agntspace.llm.base import LLMProvider, LLMResponse, Message, StreamChunk, ToolCall


class OpenAIProvider(LLMProvider):
    """LLM provider using the OpenAI Chat Completions API.

    Also works with Ollama by pointing base_url to http://localhost:11434/v1.
    """

    def __init__(
        self,
        api_key: str,
        model: str = "gpt-5.4",
        base_url: str | None = None,
    ) -> None:
        kwargs: dict = {"api_key": api_key}
        if base_url:
            kwargs["base_url"] = base_url
        self._client = openai.AsyncOpenAI(**kwargs)
        self._model = model

    async def complete(
        self,
        messages: list[Message],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        msgs = self._build_messages(messages, system)
        resp = await self._client.chat.completions.create(
            model=self._model,
            messages=msgs,
            max_tokens=max_tokens,
            temperature=temperature,
        )
        choice = resp.choices[0]
        cache_read = 0
        if resp.usage:
            details = getattr(resp.usage, "prompt_tokens_details", None)
            if details:
                cache_read = getattr(details, "cached_tokens", 0) or 0
        return LLMResponse(
            content=choice.message.content or "",
            model=resp.model or self._model,
            input_tokens=resp.usage.prompt_tokens if resp.usage else 0,
            output_tokens=resp.usage.completion_tokens if resp.usage else 0,
            cache_read_tokens=cache_read,
            stop_reason=choice.finish_reason,
        )

    async def stream(
        self,
        messages: list[Message],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> AsyncIterator[StreamChunk]:
        msgs = self._build_messages(messages, system)
        full_content = ""
        stream = await self._client.chat.completions.create(
            model=self._model,
            messages=msgs,
            max_tokens=max_tokens,
            temperature=temperature,
            stream=True,
        )
        async for chunk in stream:
            if not chunk.choices:
                continue
            delta = chunk.choices[0].delta
            if delta.content:
                full_content += delta.content
                yield StreamChunk(delta=delta.content)
            if chunk.choices[0].finish_reason:
                yield StreamChunk(
                    delta="",
                    done=True,
                    response=LLMResponse(
                        content=full_content,
                        model=self._model,
                        stop_reason=chunk.choices[0].finish_reason,
                    ),
                )

    async def complete_with_tools(
        self,
        messages: list[Message],
        tools: list[dict],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        """OpenAI native function calling."""
        import json

        msgs = self._build_messages(messages, system)

        # Convert Anthropic tool format to OpenAI format
        oai_tools = []
        for t in tools:
            oai_tools.append({
                "type": "function",
                "function": {
                    "name": t["name"],
                    "description": t.get("description", ""),
                    "parameters": t.get("input_schema", {}),
                },
            })

        resp = await self._client.chat.completions.create(
            model=self._model,
            messages=msgs,
            max_tokens=max_tokens,
            temperature=temperature,
            tools=oai_tools if oai_tools else None,
        )

        choice = resp.choices[0]
        text = choice.message.content or ""
        tool_calls = []

        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                tool_calls.append(ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    input=json.loads(tc.function.arguments),
                ))

        cache_read = 0
        if resp.usage:
            details = getattr(resp.usage, "prompt_tokens_details", None)
            if details:
                cache_read = getattr(details, "cached_tokens", 0) or 0
        return LLMResponse(
            content=text,
            model=resp.model or self._model,
            input_tokens=resp.usage.prompt_tokens if resp.usage else 0,
            output_tokens=resp.usage.completion_tokens if resp.usage else 0,
            cache_read_tokens=cache_read,
            stop_reason=choice.finish_reason,
            tool_calls=tool_calls,
            raw_content=choice.message,
        )

    @staticmethod
    def _build_messages(
        messages: list[Message], system: str | None
    ) -> list[dict]:
        msgs: list[dict] = []
        if system:
            msgs.append({"role": "system", "content": system})
        msgs.extend({"role": m.role, "content": m.content} for m in messages)
        return msgs
