"""
Token counter for Vertex AI Partner Models (Anthropic Claude, Mistral, etc.)

This handler provides token counting for partner models hosted on Vertex AI.
Unlike Gemini models which use Google's token counting API, partner models use
their respective publisher-specific count-tokens endpoints.
"""

from typing import Any, Dict, Optional

from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase


class VertexAIPartnerModelsTokenCounter(VertexBase):
    """
    Token counter for Vertex AI Partner Models.

    Handles token counting for models like Claude (Anthropic), Mistral, etc.
    that are available through Vertex AI's partner model program.
    """

    def _get_publisher_for_model(self, model: str) -> str:
        """
        Determine the publisher name for the given model.

        Args:
            model: The model name (e.g., "claude-3-5-sonnet-20241022")

        Returns:
            Publisher name to use in the Vertex AI endpoint URL

        Raises:
            ValueError: If the model is not a recognized partner model
        """
        if "claude" in model:
            return "anthropic"
        elif "mistral" in model or "codestral" in model:
            return "mistralai"
        elif "llama" in model or "meta/" in model:
            return "meta"
        else:
            raise ValueError(f"Unknown partner model: {model}")

    def _build_count_tokens_endpoint(
        self,
        model: str,
        project_id: str,
        vertex_location: str,
        api_base: Optional[str] = None,
    ) -> str:
        """
        Build the count-tokens endpoint URL for a partner model.

        Args:
            model: The model name
            project_id: Google Cloud project ID
            vertex_location: Vertex AI location (e.g., "us-east5")
            api_base: Optional custom API base URL

        Returns:
            Full endpoint URL for the count-tokens API
        """
        publisher = self._get_publisher_for_model(model)

        # Use custom api_base if provided, otherwise construct default
        if api_base:
            base_url = api_base
        else:
            base_url = get_vertex_base_url(vertex_location)

        # Construct the count-tokens endpoint
        # Format: /v1/projects/{project}/locations/{location}/publishers/{publisher}/models/count-tokens:rawPredict
        endpoint = (
            f"{base_url}/v1/projects/{project_id}/locations/{vertex_location}/"
            f"publishers/{publisher}/models/count-tokens:rawPredict"
        )

        return endpoint

    @staticmethod
    def _strip_version_suffix(model: str) -> str:
        """
        Strip version suffixes (e.g. @default, @20251001) from model names.

        The Vertex AI count-tokens endpoint rejects model names that include
        version suffixes — for example, "claude-sonnet-4-6@default" returns
        "not supported for token counting" while "claude-sonnet-4-6" works.
        """
        if "@" in model:
            return model.split("@")[0]
        return model

    async def handle_count_tokens_request(
        self,
        model: str,
        request_data: Dict[str, Any],
        litellm_params: Dict[str, Any],
    ) -> Dict[str, Any]:
        """
        Handle token counting request for a Vertex AI partner model.

        Args:
            model: The model name
            request_data: Request payload (Anthropic Messages API format)
            litellm_params: LiteLLM parameters containing credentials, project, location

        Returns:
            Dict containing token count information

        Raises:
            ValueError: If required parameters are missing or invalid
        """
        # Strip version suffixes (@default, @20251001, etc.) — the Vertex AI
        # count-tokens endpoint does not accept versioned model names.
        model = self._strip_version_suffix(model)
        if "model" in request_data:
            request_data = {
                **request_data,
                "model": self._strip_version_suffix(request_data["model"]),
            }

        # Validate request
        if "messages" not in request_data:
            raise ValueError("messages required for token counting")

        # Extract Vertex AI credentials and settings
        vertex_credentials = self.get_vertex_ai_credentials(litellm_params)
        vertex_project = self.get_vertex_ai_project(litellm_params)

        # Check for count_tokens specific location override
        vertex_count_tokens_location = litellm_params.get(
            "vertex_count_tokens_location"
        )
        vertex_location_raw = self.get_vertex_ai_location(litellm_params)

        # Determine final location with precedence:
        # 1. vertex_count_tokens_location (if provided)
        # 2. vertex_location (if provided)
        # 3. Default to us-east5 for Claude models when no location is set
        # Supported regions: us-east5, europe-west1, asia-southeast1
        # https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens
        if vertex_count_tokens_location:
            vertex_location: str = vertex_count_tokens_location
        elif vertex_location_raw:
            vertex_location = vertex_location_raw
        elif "claude" in model.lower():
            vertex_location = "us-east5"
        else:
            vertex_location = "us-east5"

        # Get access token and resolved project ID
        access_token, project_id = await self._ensure_access_token_async(
            credentials=vertex_credentials,
            project_id=vertex_project,
            custom_llm_provider="vertex_ai",
        )

        # Build the endpoint URL
        endpoint_url = self._build_count_tokens_endpoint(
            model=model,
            project_id=project_id,
            vertex_location=vertex_location,
            api_base=litellm_params.get("api_base"),
        )

        # Prepare headers
        headers = {"Authorization": f"Bearer {access_token}"}

        # Get async HTTP client
        from litellm import LlmProviders

        async_client = get_async_httpx_client(llm_provider=LlmProviders.VERTEX_AI)

        # Make the request
        # Note: Partner models (especially Claude) accept Anthropic Messages API format directly
        response = await async_client.post(
            endpoint_url,
            headers=headers,
            json=request_data,
            timeout=30.0,
        )

        # Check for errors
        if response.status_code != 200:
            error_text = response.text
            raise ValueError(
                f"Token counting request failed with status {response.status_code}: {error_text}"
            )

        # Parse response
        result = response.json()

        # Return token count
        # Vertex AI Anthropic returns: {"input_tokens": 123}
        return {
            "input_tokens": result.get("input_tokens", 0),
            "tokenizer_used": "vertex_ai_partner_models",
        }
