"""
Transformation for NVIDIA NIM Ranking models that use /v1/ranking endpoint.

Use this by passing "nvidia_nim/ranking/<model>" to force the /v1/ranking endpoint.

Reference: https://build.nvidia.com/nvidia/llama-3_2-nv-rerankqa-1b-v2/deploy
"""

from typing import Dict, Optional

from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig


class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
    """
    Configuration for NVIDIA NIM models that use the /v1/ranking endpoint.
    
    Example:
        curl -X "POST" 'https://ai.api.nvidia.com/v1/ranking' \
            -H 'Accept: application/json' \
            -H 'Content-Type: application/json' \
            -d '{
                "model": "nvidia/llama-3.2-nv-rerankqa-1b-v2",
                "query": {"text": "which way did the traveler go?"},
                "passages": [{"text": "..."}, {"text": "..."}],
                "truncate": "END"
            }'
    """

    def _get_clean_model_name(self, model: str) -> str:
        """Strip 'nvidia_nim/' and 'ranking/' prefixes from model name."""
        # First strip nvidia_nim/ prefix if present
        if model.startswith("nvidia_nim/"):
            model = model[len("nvidia_nim/"):]
        # Then strip ranking/ prefix if present
        if model.startswith("ranking/"):
            model = model[len("ranking/"):]
        return model

    def get_complete_url(
        self,
        api_base: Optional[str],
        model: str,
        optional_params: Optional[dict] = None,
    ) -> str:
        """
        Construct the Nvidia NIM ranking URL.
        
        Format: {api_base}/v1/ranking
        """
        if not api_base:
            api_base = self.DEFAULT_NIM_RERANK_API_BASE

        api_base = api_base.rstrip("/")

        if api_base.endswith("/ranking"):
            return api_base

        if api_base.endswith("/v1"):
            api_base = api_base[:-3]

        return f"{api_base}/v1/ranking"

    def transform_rerank_request(
        self,
        model: str,
        optional_rerank_params: Dict,
        headers: dict,
    ) -> dict:
        """
        Transform request, using clean model name without 'ranking/' prefix.
        """
        clean_model = self._get_clean_model_name(model)
        return super().transform_rerank_request(
            model=clean_model,
            optional_rerank_params=optional_rerank_params,
            headers=headers,
        )