"""OpenAI provider implementation."""

from __future__ import annotations

import json
from typing import TYPE_CHECKING

from koan.llm.retry import llm_retry
from koan.types import LLMResponse, Message, StreamDelta, Usage

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.llm.config import LLMConfig


def _serialize_messages(messages: list[Message]) -> list[dict]:
    """Convert KOAN Messages to OpenAI message dicts, preserving tool_calls and tool_call_id."""
    result = []
    for m in messages:
        msg: dict = {"role": m.role, "content": m.content}
        if m.role == "assistant" and m.tool_calls:
            msg["tool_calls"] = [
                {
                    "id": tc["id"],
                    "type": "function",
                    "function": {
                        "name": tc["name"],
                        "arguments": (
                            tc["arguments"] if isinstance(tc["arguments"], str)
                            else json.dumps(tc["arguments"])
                        ),
                    },
                }
                for tc in m.tool_calls
            ]
        if m.role == "tool" and m.tool_call_id:
            msg["tool_call_id"] = m.tool_call_id
        result.append(msg)
    return result


class OpenAIProvider:
    """Async OpenAI completions provider."""

    def __init__(self, api_key: str) -> None:
        try:
            from openai import AsyncOpenAI
        except ImportError as e:
            raise ImportError(
                "OpenAI SDK not installed. Install with: pip install koan[openai]"
            ) from e
        self._client = AsyncOpenAI(api_key=api_key)

    @llm_retry
    async def complete(self, config: LLMConfig, messages: list[Message]) -> LLMResponse:
        """Call OpenAI Chat Completions API and return a normalized response."""
        kwargs: dict = {
            "model": config.model,
            "messages": _serialize_messages(messages),
            "temperature": config.temperature,
            "max_tokens": config.max_tokens,
            "top_p": config.top_p,
            "timeout": config.timeout,
        }
        if config.tools:
            kwargs["tools"] = config.tools

        response = await self._client.chat.completions.create(**kwargs)  # type: ignore[arg-type]

        choice = response.choices[0]
        usage = response.usage

        # Extract tool calls into raw dict for agent consumption
        raw: dict | None = None
        if choice.message.tool_calls:
            raw = {
                "tool_calls": [
                    {
                        "id": tc.id,
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    }
                    for tc in choice.message.tool_calls
                ],
            }

        return LLMResponse(
            text=choice.message.content or "",
            model=response.model,
            provider="openai",
            usage=Usage(
                input_tokens=usage.prompt_tokens if usage else 0,
                output_tokens=usage.completion_tokens if usage else 0,
            ),
            stop_reason=choice.finish_reason,
            raw=raw,
        )

    async def stream(self, config: LLMConfig, messages: list[Message]) -> AsyncIterator[StreamDelta]:
        """Stream OpenAI Chat Completions and yield normalized deltas."""
        kwargs: dict = {
            "model": config.model,
            "messages": _serialize_messages(messages),
            "temperature": config.temperature,
            "max_tokens": config.max_tokens,
            "top_p": config.top_p,
            "timeout": config.timeout,
            "stream": True,
            "stream_options": {"include_usage": True},
        }
        if config.tools:
            kwargs["tools"] = config.tools

        response = await self._client.chat.completions.create(**kwargs)  # type: ignore[arg-type]

        async for chunk in response:
            if not chunk.choices:
                # Final chunk with usage only
                if chunk.usage:
                    yield StreamDelta(
                        provider="openai",
                        model=chunk.model,
                        usage=Usage(
                            input_tokens=chunk.usage.prompt_tokens,
                            output_tokens=chunk.usage.completion_tokens,
                        ),
                    )
                continue

            choice = chunk.choices[0]
            delta = choice.delta

            # Tool call deltas
            if delta.tool_calls:
                for tc in delta.tool_calls:
                    yield StreamDelta(
                        provider="openai",
                        model=chunk.model,
                        tool_call_id=tc.id or None,
                        tool_name=tc.function.name if tc.function and tc.function.name else None,
                        tool_args_delta=tc.function.arguments if tc.function and tc.function.arguments else None,
                        finish_reason=choice.finish_reason,
                    )
            else:
                yield StreamDelta(
                    text=delta.content or "",
                    provider="openai",
                    model=chunk.model,
                    finish_reason=choice.finish_reason,
                )

    async def close(self) -> None:
        await self._client.close()
