"""OpenAI (GPT) provider client."""

import json
from collections.abc import AsyncGenerator
from typing import Any

import structlog

from robot_resources.proxy.models import (
    ChatCompletionRequest,
    ChatCompletionResponse,
    Choice,
    Message,
    Usage,
)
from robot_resources.proxy.providers.base import BaseProvider

logger = structlog.get_logger()


class OpenAIProvider(BaseProvider):
    """OpenAI GPT API provider."""

    @property
    def name(self) -> str:
        return "openai"

    @property
    def base_url(self) -> str:
        return "https://api.openai.com"

    def _get_headers(self, *, api_key: str | None = None) -> dict:
        key = api_key or self.config.api_key
        return {
            "Authorization": f"Bearer {key}",
            "content-type": "application/json",
        }

    def _get_completion_endpoint(self) -> str:
        return "/v1/chat/completions"

    def map_model(self, model: str) -> str:
        """Map model name (OpenAI uses standard names)."""
        return model

    def format_request(self, request: ChatCompletionRequest) -> dict[str, Any]:
        """Format request for OpenAI API (mostly passthrough)."""
        formatted: dict[str, Any] = {
            "model": self.map_model(request.model),
            "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
        }

        if request.max_tokens:
            # gpt-5.4+ models require max_completion_tokens instead of max_tokens
            formatted["max_completion_tokens"] = request.max_tokens

        if request.temperature != 1.0:
            formatted["temperature"] = request.temperature

        return formatted

    def parse_response(self, response: dict) -> ChatCompletionResponse:
        """Parse OpenAI response to standard format (mostly passthrough)."""
        choice = response["choices"][0]
        usage = response.get("usage", {})

        return ChatCompletionResponse(
            id=response.get("id", ""),
            model=response.get("model", ""),
            choices=[
                Choice(
                    index=choice.get("index", 0),
                    message=Message(
                        role=choice["message"]["role"],
                        content=choice["message"]["content"],
                    ),
                    finish_reason=choice.get("finish_reason", "stop"),
                )
            ],
            usage=Usage(
                prompt_tokens=usage.get("prompt_tokens", 0),
                completion_tokens=usage.get("completion_tokens", 0),
                total_tokens=usage.get("total_tokens", 0),
            ),
        )

    async def stream(
        self, request: ChatCompletionRequest, *, api_key: str | None = None,
    ) -> AsyncGenerator[str, None]:
        """Stream a chat completion from OpenAI.

        OpenAI's streaming format is already ChatCompletionChunk-compatible,
        so this is near-passthrough.  We request ``include_usage`` so the
        final chunk contains token counts for tracking.
        """
        formatted = self.format_request(request)
        formatted["stream"] = True
        formatted["stream_options"] = {"include_usage": True}

        url = f"{self.config.base_url or self.base_url}{self._get_completion_endpoint()}"
        self._last_stream_usage = None

        async with self._create_client(api_key=api_key) as client:
            async with client.stream("POST", url, json=formatted) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    line = line.strip()
                    if not line:
                        continue
                    if not line.startswith("data: "):
                        continue

                    payload = line[len("data: "):]

                    if payload == "[DONE]":
                        yield "data: [DONE]\n\n"
                        return

                    # Extract usage from the final chunk if present
                    try:
                        chunk_data = json.loads(payload)
                        usage = chunk_data.get("usage")
                        if usage:
                            self._last_stream_usage = Usage(
                                prompt_tokens=usage.get("prompt_tokens", 0),
                                completion_tokens=usage.get("completion_tokens", 0),
                                total_tokens=usage.get("total_tokens", 0),
                            )
                    except json.JSONDecodeError:
                        logger.warning(
                            "stream_json_decode_error",
                            provider="openai",
                            raw_line=payload[:200],
                        )

                    yield f"data: {payload}\n\n"
