from typing import List, Optional, Union
from ..transport import SyncTransport
from ..models.rerank import RerankResponse


class RerankResource:
    def __init__(self, transport: SyncTransport):
        self._transport = transport

    def create(
        self,
        query: str,
        documents: List[Union[str, dict]],
        *,
        model: str,
        top_n: Optional[int] = None,
        return_documents: bool = True,
    ) -> RerankResponse:
        """Rerank documents by relevance to a query.

        Args:
            query: The search query.
            documents: List of documents (strings or dicts) to rerank.
            model: Reranking model name.
            top_n: Number of top results to return (None returns all).
            return_documents: Whether to include document text in results.

        Returns:
            RerankResponse with ranked results.
        """
        payload = {
            "query": query,
            "documents": documents,
            "model": model,
            "return_documents": return_documents,
        }
        if top_n is not None:
            payload["top_n"] = top_n
        data = self._transport.request("POST", "/rerank", json=payload)
        return RerankResponse(**data)
