"""Inference response protocols."""

from typing import Any, Protocol, runtime_checkable

from pinecone.__interface__.inference.results import RerankResultProtocol
from pinecone.__interface__.inference.usage import UsageProtocol


@runtime_checkable
class EmbeddingResponseProtocol(Protocol):
    """Protocol for embedding generation responses.

    Example:
        >>> response = pc.inference.embed(model="multilingual-e5-large", texts=["Hello world"])
        >>> embeddings = response.embeddings
        >>> print(f"Tokens used: {response.usage.total_tokens}")
    """

    @property
    def embeddings(self) -> list[list[float]]:
        """Generated embeddings."""
        ...

    @property
    def model(self) -> str:
        """Model name used."""
        ...

    @property
    def usage(self) -> UsageProtocol:
        """Token usage information."""
        ...


@runtime_checkable
class RerankResponseProtocol(Protocol):
    """Protocol for reranking operation responses.

    Example:
        >>> response = pc.inference.rerank(
        ...     model="bge-reranker-v2-m3",
        ...     query="What is deep learning?",
        ...     documents=["Deep learning is...", "Machine learning is..."],
        ... )
        >>> for result in response.results:
        ...     print(f"Score: {result.score}, Index: {result.index}")
        >>> print(f"Tokens used: {response.usage.total_tokens}")
    """

    @property
    def results(self) -> list[RerankResultProtocol]:
        """Reranking results with scores and indices."""
        ...

    @property
    def model(self) -> str:
        """Model name used."""
        ...

    @property
    def usage(self) -> UsageProtocol:
        """Token usage information."""
        ...


@runtime_checkable
class ListModelsResponseProtocol(Protocol):
    """Protocol for list_models operation responses.

    Example:
        >>> response = pc.inference.list_models()
        >>> for model in response.models:
        ...     print(f"{model.model}: {model.short_description}")
    """

    @property
    def models(self) -> list[Any]:
        """List of available models."""
        ...
