"""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 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
