"""
OpenAI provider adapter.

Since the proxy speaks OpenAI-compatible format natively, this adapter
is mostly a passthrough. It handles auth header formatting and pricing lookup.
"""

from __future__ import annotations

import json
from typing import Any, Dict, Tuple

from proxy.adapters.base import BaseProviderAdapter
from proxy.pricing import OPENAI_PRICING, ZERO_PRICING
from proxy.schemas import (
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatCompletionResponse,
    Choice,
    ChoiceMessage,
    StreamChoice,
    StreamDelta,
    TokenPricing,
    Usage,
)


class OpenAIAdapter(BaseProviderAdapter):
    """Adapter for the OpenAI API.

    The proxy's native format IS OpenAI-compatible, so translation is
    minimal — mainly auth header and response parsing.
    """

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

    # ----- request translation -----

    def translate_request(
        self,
        request: ChatCompletionRequest,
        api_key: str,
    ) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
        url = f"{self.base_url}/v1/chat/completions"

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
        }

        body: Dict[str, Any] = {
            "model": request.model,
            "messages": [m.model_dump(exclude_none=True) for m in request.messages],
        }
        if request.temperature is not None:
            body["temperature"] = request.temperature
        if request.top_p is not None:
            body["top_p"] = request.top_p
        if request.n is not None:
            body["n"] = request.n
        if request.stream:
            body["stream"] = True
            body["stream_options"] = {"include_usage": True}
        if request.stop is not None:
            body["stop"] = request.stop
        if request.max_tokens is not None:
            body["max_tokens"] = request.max_tokens
        if request.max_completion_tokens is not None:
            body["max_completion_tokens"] = request.max_completion_tokens
        if request.presence_penalty:
            body["presence_penalty"] = request.presence_penalty
        if request.frequency_penalty:
            body["frequency_penalty"] = request.frequency_penalty
        if request.logit_bias:
            body["logit_bias"] = request.logit_bias
        if request.user:
            body["user"] = request.user

        return url, headers, body

    # ----- response translation -----

    def translate_response(
        self,
        status_code: int,
        headers: Dict[str, str],
        body: Dict[str, Any],
    ) -> ChatCompletionResponse:
        # OpenAI response is already in our unified format
        usage_data = body.get("usage", {})
        choices_data = body.get("choices", [])

        choices = []
        for c in choices_data:
            msg = c.get("message", {})
            choices.append(
                Choice(
                    index=c.get("index", 0),
                    message=ChoiceMessage(
                        role=msg.get("role", "assistant"),
                        content=msg.get("content"),
                        tool_calls=msg.get("tool_calls"),
                    ),
                    finish_reason=c.get("finish_reason"),
                )
            )

        return ChatCompletionResponse(
            id=body.get("id", ""),
            object=body.get("object", "chat.completion"),
            created=body.get("created", 0),
            model=body.get("model", ""),
            choices=choices,
            usage=Usage(
                prompt_tokens=usage_data.get("prompt_tokens", 0),
                completion_tokens=usage_data.get("completion_tokens", 0),
                total_tokens=usage_data.get("total_tokens", 0),
            ),
        )

    # ----- streaming -----

    def parse_stream_chunk(self, raw_line: str) -> ChatCompletionChunk | None:
        line = raw_line.strip()
        if not line or line == "[DONE]":
            return None

        try:
            data = json.loads(line)
        except json.JSONDecodeError:
            return None

        choices_data = data.get("choices", [])
        stream_choices = []
        for c in choices_data:
            delta = c.get("delta", {})
            stream_choices.append(
                StreamChoice(
                    index=c.get("index", 0),
                    delta=StreamDelta(
                        role=delta.get("role"),
                        content=delta.get("content"),
                    ),
                    finish_reason=c.get("finish_reason"),
                )
            )

        return ChatCompletionChunk(
            id=data.get("id", ""),
            object=data.get("object", "chat.completion.chunk"),
            created=data.get("created", 0),
            model=data.get("model", ""),
            choices=stream_choices,
        )

    # ----- pricing -----

    def get_pricing(self, model: str) -> TokenPricing:
        return OPENAI_PRICING.get(model, ZERO_PRICING)
