"""Data models for document reranking API.

API Version: 2025-10
"""

from typing import TYPE_CHECKING, Any

from pinecone.models._base import BaseModel
from pinecone.models._display import render_table
from pinecone.models.inference._common import Usage


if TYPE_CHECKING:
    from pinecone.__interface__.inference import RerankResultProtocol


class RerankDocument(BaseModel, kw_only=True):  # pyright: ignore[reportCallIssue]
    """Document to be reranked.

    Attributes:
        text: The document text
        id: Optional document identifier
    """

    text: str
    id: str | None = None


class RerankResult(BaseModel, kw_only=True):  # pyright: ignore[reportCallIssue]
    """Single reranking result.

    Attributes:
        index: Index of the document in the original input list
        score: Relevance score (higher is more relevant)
        document: Optional document text (if return_documents=True)
    """

    index: int
    score: float
    document: str | None = None


class RerankRequest(BaseModel, kw_only=True):  # pyright: ignore[reportCallIssue]
    """Request model for document reranking.

    Attributes:
        model: Name of the reranking model to use
        query: The query text
        documents: List of documents to rerank
        top_n: Optional number of top results to return
        return_documents: Whether to return document text in results
        parameters: Optional model-specific parameters
    """

    model: str
    query: str
    documents: list[RerankDocument]
    top_n: int | None = None
    return_documents: bool = False
    parameters: dict[str, Any] | None = None


class RerankResponse(BaseModel, kw_only=True):  # pyright: ignore[reportCallIssue]
    """Response model for document reranking.

    Attributes:
        data: List of reranking results
        model: Name of the model used
        usage: Token usage information
    """

    data: list[RerankResult]
    model: str
    usage: Usage

    @property
    def results(self) -> "list[RerankResultProtocol]":
        """Get reranking results with full type information.

        Returns:
            List of RerankResult objects with index, score, and optional document.

        Example:
            >>> response = pc.inference.rerank(...)
            >>> for result in response.results:
            ...     print(f"Index: {result.index}, Score: {result.score}")
        """
        return self.data  # type: ignore[return-value]

    def __repr__(self) -> str:
        """Return summary representation without overwhelming array data."""
        return (
            f"RerankResponse(model={self.model!r}, results={len(self.data)}, usage={self.usage!r})"
        )

    def __dir__(self) -> list[str]:
        """Return public API."""
        attrs = set(super().__dir__())
        public = {name for name in attrs if not name.startswith("_")}
        return sorted(public)

    def _repr_pretty_(self, p: Any, cycle: bool) -> None:
        """IPython pretty-printer support."""
        if cycle:
            p.text("RerankResponse(...)")
            return

        p.text("RerankResponse(")
        with p.group(2, "", ")"):
            p.breakable()
            p.text(f"model={self.model!r},")
            p.breakable()
            p.text(f"results={len(self.data)},")
            p.breakable()
            p.text(f"usage={self.usage!r}")

    def _repr_html_(self) -> str:
        """Jupyter notebook HTML representation."""
        return render_table(
            "RerankResponse",
            [
                ("Model:", self.model),
                ("Results:", len(self.data)),
                ("Total Tokens:", self.usage.total_tokens),
            ],
        )
