"""Adapter for transforming Inference API responses to SDK models.

This adapter provides minimal transformation from raw API responses to typed SDK models.
All transformations are documented to explain why they're needed.
"""

from typing import Any

from pinecone.models.inference import (
    Embedding,
    EmbedResponse,
    ListModelsResponse,
    ModelInfo,
    ModelInfoSupportedParameter,
    RerankResponse,
    RerankResult,
    Usage,
)


class InferenceAdapter:
    """Adapter for Inference API responses.

    This class transforms raw API responses into typed SDK models. Each transformation
    is kept minimal and documented.
    """

    @staticmethod
    def to_embed_response(data: dict[str, Any]) -> EmbedResponse:
        """Transform embed API response to EmbedResponse model.

        Args:
            data: Raw API response dictionary

        Returns:
            Typed EmbedResponse model

        Transformations:
            - Wraps embedding values in Embedding objects with indices
            - Extracts usage information into Usage model
            - Preserves model name from API response
        """
        embeddings = [
            Embedding(values=item["values"], index=i) for i, item in enumerate(data["data"])
        ]

        return EmbedResponse(
            data=embeddings,
            model=data["model"],
            vector_type=data.get("vector_type", "dense"),  # Default to dense for backwards compat
            usage=Usage(total_tokens=data["usage"]["total_tokens"]),
        )

    @staticmethod
    def to_rerank_response(data: dict[str, Any]) -> RerankResponse:
        """Transform rerank API response to RerankResponse model.

        Args:
            data: Raw API response dictionary

        Returns:
            Typed RerankResponse model

        Transformations:
            - Wraps reranking results in RerankResult objects
            - Preserves index, score, and optional document text
            - Extracts usage information into Usage model
            - Preserves model name from API response
        """
        results = [
            RerankResult(
                index=item["index"],
                score=item["score"],
                document=item.get("document"),  # Optional field
            )
            for item in data["data"]
        ]

        # Rerank API returns "rerank_units" not "total_tokens"
        usage_data = data["usage"]
        total_tokens = usage_data.get("total_tokens") or usage_data.get("rerank_units", 0)

        return RerankResponse(
            data=results,
            model=data["model"],
            usage=Usage(total_tokens=total_tokens),
        )

    @staticmethod
    def to_list_models_response(data: dict[str, Any]) -> ListModelsResponse:
        """Transform list models API response to ListModelsResponse model.

        Args:
            data: Raw API response dictionary

        Returns:
            Typed ListModelsResponse model

        Transformations:
            - Wraps model information in ModelInfo objects
            - Wraps supported parameters in ModelInfoSupportedParameter objects
            - Preserves all model metadata from API response
        """
        models = []
        for item in data["models"]:
            # Parse supported parameters
            supported_params = [
                ModelInfoSupportedParameter(
                    parameter=param["parameter"],
                    type=param["type"],
                    value_type=param["value_type"],
                    required=param["required"],
                    allowed_values=param.get("allowed_values"),
                    min=param.get("min"),
                    max=param.get("max"),
                    default=param.get("default"),
                )
                for param in item["supported_parameters"]
            ]

            models.append(
                ModelInfo(
                    model=item["model"],
                    short_description=item["short_description"],
                    type=item["type"],
                    modality=item["modality"],
                    max_sequence_length=item["max_sequence_length"],
                    max_batch_size=item["max_batch_size"],
                    provider_name=item["provider_name"],
                    supported_parameters=supported_params,
                    vector_type=item.get("vector_type"),
                    default_dimension=item.get("default_dimension"),
                    supported_dimensions=item.get("supported_dimensions"),
                    supported_metrics=item.get("supported_metrics"),
                )
            )

        return ListModelsResponse(models=models)
