"""OpenAI-compatible chat, embedding, and moderation protocol."""

from typing import Any

from agenteval.models import ModelCapability, ModelProtocol
from agenteval.utils.llm import extract_content_from_llm_response

from .base import ModelProtocolAdapter, ProtocolAdapterError


class OpenAICompatibleAdapter(ModelProtocolAdapter):
    protocol = ModelProtocol.OPENAI_COMPATIBLE
    supported_capabilities = frozenset(ModelCapability)

    def headers(self, api_key: str | None) -> dict[str, str]:
        headers = super().headers(api_key)
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"
        return headers

    def chat_payload(
        self,
        model_name: str | None,
        messages: list[dict[str, str]],
        temperature: float,
    ) -> dict[str, Any]:
        return {"model": self.require_model(model_name), "messages": messages, "temperature": temperature}

    def parse_chat(self, data: dict[str, Any]) -> str:
        content = extract_content_from_llm_response(data)
        if not content:
            raise ProtocolAdapterError("模型接口返回内容为空")
        return content

    def parse_usage(self, data: dict[str, Any]) -> dict[str, int] | None:
        """Extract token usage from OpenAI-compatible response."""
        usage = data.get("usage")
        if not usage or not isinstance(usage, dict):
            return None
        return {
            "prompt_tokens": usage.get("prompt_tokens", 0),
            "completion_tokens": usage.get("completion_tokens", 0),
            "total_tokens": usage.get("total_tokens", 0),
        }

    def embedding_payload(self, model_name: str | None, inputs: str | list[str]) -> dict[str, Any]:
        return {"model": self.require_model(model_name), "input": inputs}

    def parse_embeddings(self, data: dict[str, Any]) -> list[list[float]]:
        try:
            rows = sorted(data["data"], key=lambda item: item.get("index", 0))
            vectors = [row["embedding"] for row in rows]
        except (KeyError, TypeError) as exc:
            raise ProtocolAdapterError("Embedding 接口返回格式不正确") from exc
        if not all(isinstance(vector, list) for vector in vectors):
            raise ProtocolAdapterError("Embedding 接口返回格式不正确")
        return vectors

    def moderation_payload(self, model_name: str | None, text: str) -> dict[str, Any]:
        payload: dict[str, Any] = {"input": text}
        if model_name:
            payload["model"] = model_name
        return payload

    def parse_moderation(self, data: dict[str, Any]) -> dict[str, Any]:
        try:
            result = data["results"][0]
        except (KeyError, IndexError, TypeError) as exc:
            raise ProtocolAdapterError("Moderation 接口返回格式不正确") from exc
        if not isinstance(result, dict):
            raise ProtocolAdapterError("Moderation 接口返回格式不正确")
        return result