"""
Batch Rate Limiter Hook

This hook implements rate limiting for batch API requests by:
1. Reading batch input files to count requests and estimate tokens at submission
2. Validating actual usage from output files when batches complete
3. Integrating with the existing parallel request limiter infrastructure

## Integration & Calling
This hook is automatically registered and called by the proxy system.
See BATCH_RATE_LIMITER_INTEGRATION.md for complete integration details.

Quick summary:
- Add to PROXY_HOOKS in litellm/proxy/hooks/__init__.py
- Gets auto-instantiated on proxy startup via _add_proxy_hooks()
- async_pre_call_hook() fires on POST /v1/batches (batch submission)
- async_log_success_event() fires on GET /v1/batches/{id} (batch completion)
"""

from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union

from fastapi import HTTPException
from pydantic import BaseModel

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import (
    _get_batch_job_input_file_usage,
    _get_file_content_as_dictionary,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth

if TYPE_CHECKING:
    from opentelemetry.trace import Span as _Span

    from litellm.proxy.hooks.parallel_request_limiter_v3 import (
        RateLimitDescriptor as _RateLimitDescriptor,
    )
    from litellm.proxy.hooks.parallel_request_limiter_v3 import (
        RateLimitStatus as _RateLimitStatus,
    )
    from litellm.proxy.hooks.parallel_request_limiter_v3 import (
        _PROXY_MaxParallelRequestsHandler_v3 as _ParallelRequestLimiter,
    )
    from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
    from litellm.router import Router as _Router

    Span = Union[_Span, Any]
    InternalUsageCache = _InternalUsageCache
    Router = _Router
    ParallelRequestLimiter = _ParallelRequestLimiter
    RateLimitStatus = _RateLimitStatus
    RateLimitDescriptor = _RateLimitDescriptor
else:
    Span = Any
    InternalUsageCache = Any
    Router = Any
    ParallelRequestLimiter = Any
    RateLimitStatus = Dict[str, Any]
    RateLimitDescriptor = Dict[str, Any]


class BatchFileUsage(BaseModel):
    """
    Internal model for batch file usage tracking, used for batch rate limiting
    """

    total_tokens: int
    request_count: int


class _PROXY_BatchRateLimiter(CustomLogger):
    """
    Rate limiter for batch API requests.

    Handles rate limiting at two points:
    1. Batch submission - reads input file and reserves capacity
    2. Batch completion - reads output file and adjusts for actual usage
    """

    def __init__(
        self,
        internal_usage_cache: InternalUsageCache,
        parallel_request_limiter: ParallelRequestLimiter,
    ):
        """
        Initialize the batch rate limiter.

        Note: These dependencies are automatically injected by ProxyLogging._add_proxy_hooks()
        when this hook is registered in PROXY_HOOKS. See BATCH_RATE_LIMITER_INTEGRATION.md.

        Args:
            internal_usage_cache: Cache for storing rate limit data (auto-injected)
            parallel_request_limiter: Existing rate limiter to integrate with (needs custom injection)
        """
        self.internal_usage_cache = internal_usage_cache
        self.parallel_request_limiter = parallel_request_limiter

    def _raise_rate_limit_error(
        self,
        status: "RateLimitStatus",
        descriptors: List["RateLimitDescriptor"],
        batch_usage: BatchFileUsage,
        limit_type: str,
    ) -> None:
        """Raise HTTPException for rate limit exceeded."""
        from datetime import datetime

        # Find the descriptor for this status
        descriptor_index = next(
            (
                i
                for i, d in enumerate(descriptors)
                if d.get("key") == status.get("descriptor_key")
            ),
            0,
        )
        descriptor: RateLimitDescriptor = (
            descriptors[descriptor_index]
            if descriptors
            else {"key": "", "value": "", "rate_limit": None}
        )

        now = datetime.now().timestamp()
        window_size = self.parallel_request_limiter.window_size
        reset_time = now + window_size
        reset_time_formatted = datetime.fromtimestamp(reset_time).strftime(
            "%Y-%m-%d %H:%M:%S UTC"
        )

        remaining_display = max(0, status["limit_remaining"])
        current_limit = status["current_limit"]

        if limit_type == "requests":
            detail = (
                f"Batch rate limit exceeded for {descriptor.get('key', 'unknown')}: {descriptor.get('value', 'unknown')}. "
                f"Batch contains {batch_usage.request_count} requests but only {remaining_display} requests remaining "
                f"out of {current_limit} RPM limit. "
                f"Limit resets at: {reset_time_formatted}"
            )
        else:  # tokens
            detail = (
                f"Batch rate limit exceeded for {descriptor.get('key', 'unknown')}: {descriptor.get('value', 'unknown')}. "
                f"Batch contains {batch_usage.total_tokens} tokens but only {remaining_display} tokens remaining "
                f"out of {current_limit} TPM limit. "
                f"Limit resets at: {reset_time_formatted}"
            )

        raise HTTPException(
            status_code=429,
            detail=detail,
            headers={
                "retry-after": str(window_size),
                "rate_limit_type": limit_type,
                "reset_at": reset_time_formatted,
            },
        )

    async def _check_and_increment_batch_counters(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        data: Dict,
        batch_usage: BatchFileUsage,
    ) -> None:
        """
        Check rate limits and increment counters by the batch amounts.

        Raises HTTPException if any limit would be exceeded.
        """
        from litellm.types.caching import RedisPipelineIncrementOperation

        # Create descriptors and check if batch would exceed limits
        descriptors = self.parallel_request_limiter._create_rate_limit_descriptors(
            user_api_key_dict=user_api_key_dict,
            data=data,
            rpm_limit_type=None,
            tpm_limit_type=None,
            model_has_failures=False,
        )

        # Check current usage without incrementing
        rate_limit_response = await self.parallel_request_limiter.should_rate_limit(
            descriptors=descriptors,
            parent_otel_span=user_api_key_dict.parent_otel_span,
            read_only=True,
        )

        # Verify batch won't exceed any limits
        for status in rate_limit_response["statuses"]:
            rate_limit_type = status["rate_limit_type"]
            limit_remaining = status["limit_remaining"]

            required_capacity = (
                batch_usage.request_count
                if rate_limit_type == "requests"
                else batch_usage.total_tokens
                if rate_limit_type == "tokens"
                else 0
            )

            if required_capacity > limit_remaining:
                self._raise_rate_limit_error(
                    status, descriptors, batch_usage, rate_limit_type
                )

        # Build pipeline operations for batch increments
        # Reuse the same keys that descriptors check
        pipeline_operations: List[RedisPipelineIncrementOperation] = []

        for descriptor in descriptors:
            key = descriptor["key"]
            value = descriptor["value"]
            rate_limit = descriptor.get("rate_limit")

            if rate_limit is None:
                continue

            # Add RPM increment if limit is set
            if rate_limit.get("requests_per_unit") is not None:
                rpm_key = self.parallel_request_limiter.create_rate_limit_keys(
                    key=key, value=value, rate_limit_type="requests"
                )
                pipeline_operations.append(
                    RedisPipelineIncrementOperation(
                        key=rpm_key,
                        increment_value=batch_usage.request_count,
                        ttl=self.parallel_request_limiter.window_size,
                    )
                )

            # Add TPM increment if limit is set
            if rate_limit.get("tokens_per_unit") is not None:
                tpm_key = self.parallel_request_limiter.create_rate_limit_keys(
                    key=key, value=value, rate_limit_type="tokens"
                )
                pipeline_operations.append(
                    RedisPipelineIncrementOperation(
                        key=tpm_key,
                        increment_value=batch_usage.total_tokens,
                        ttl=self.parallel_request_limiter.window_size,
                    )
                )

        # Execute increments
        if pipeline_operations:
            await self.parallel_request_limiter.async_increment_tokens_with_ttl_preservation(
                pipeline_operations=pipeline_operations,
                parent_otel_span=user_api_key_dict.parent_otel_span,
            )

    async def count_input_file_usage(
        self,
        file_id: str,
        custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
        user_api_key_dict: Optional[UserAPIKeyAuth] = None,
    ) -> BatchFileUsage:
        """
        Count number of requests and tokens in a batch input file.

        Args:
            file_id: The file ID to read
            custom_llm_provider: The custom LLM provider to use for token encoding
            user_api_key_dict: User authentication information for file access (required for managed files)

        Returns:
            BatchFileUsage with total_tokens and request_count
        """
        try:
            # Check if this is a managed file (base64 encoded unified file ID)
            from litellm.proxy.openai_files_endpoints.common_utils import (
                _is_base64_encoded_unified_file_id,
            )

            # Managed files require bypassing the HTTP endpoint (which runs access-check hooks)
            # and calling the managed files hook directly with the user's credentials.
            is_managed_file = _is_base64_encoded_unified_file_id(file_id)
            if is_managed_file and user_api_key_dict is not None:
                file_content = await self._fetch_managed_file_content(
                    file_id=file_id,
                    user_api_key_dict=user_api_key_dict,
                )
            else:
                # For non-managed files, use the standard litellm.afile_content
                file_content = await litellm.afile_content(
                    file_id=file_id,
                    custom_llm_provider=custom_llm_provider,
                    user_api_key_dict=user_api_key_dict,
                )

            file_content_as_dict = _get_file_content_as_dictionary(file_content.content)

            input_file_usage = _get_batch_job_input_file_usage(
                file_content_dictionary=file_content_as_dict,
                custom_llm_provider=custom_llm_provider,
            )
            request_count = len(file_content_as_dict)
            return BatchFileUsage(
                total_tokens=input_file_usage.total_tokens,
                request_count=request_count,
            )

        except Exception as e:
            verbose_proxy_logger.error(
                f"Error counting input file usage for {file_id}: {str(e)}"
            )
            raise

    async def _fetch_managed_file_content(
        self,
        file_id: str,
        user_api_key_dict: UserAPIKeyAuth,
    ) -> Any:
        """
        Fetch file content from managed files hook.

        This is needed for managed files because they require proper user context
        to verify file ownership and access permissions.

        Args:
            file_id: The managed file ID (base64 encoded)
            user_api_key_dict: User authentication information

        Returns:
            HttpxBinaryResponseContent with the file content
        """
        from litellm.llms.base_llm.files.transformation import BaseFileEndpoints

        # Import proxy_server dependencies at runtime to avoid circular imports
        try:
            from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
        except ImportError as e:
            raise ValueError(
                f"Cannot import proxy_server dependencies: {str(e)}. "
                "Managed files require proxy_server to be initialized."
            )

        # Get the managed files hook
        if proxy_logging_obj is None:
            raise ValueError(
                "proxy_logging_obj not available. Cannot access managed files hook."
            )

        managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files")
        if managed_files_obj is None:
            raise ValueError(
                "Managed files hook not found. Cannot access managed file."
            )

        if not isinstance(managed_files_obj, BaseFileEndpoints):
            raise ValueError("Managed files hook is not a BaseFileEndpoints instance.")

        if llm_router is None:
            raise ValueError("llm_router not available. Cannot access managed files.")

        # Use the managed files hook to get file content
        # This properly handles user permissions and file ownership
        file_content = await managed_files_obj.afile_content(
            file_id=file_id,
            litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
            llm_router=llm_router,
        )

        return file_content

    async def async_pre_call_hook(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        cache: Any,
        data: Dict,
        call_type: str,
    ) -> Union[Exception, str, Dict, None]:
        """
        Pre-call hook for batch operations.

        Only handles batch creation (acreate_batch):
        - Reads input file
        - Counts tokens and requests
        - Reserves rate limit capacity via parallel_request_limiter

        Args:
            user_api_key_dict: User authentication information
            cache: Cache instance (not used directly)
            data: Request data
            call_type: Type of call being made

        Returns:
            Modified data dict or None

        Raises:
            HTTPException: 429 if rate limit would be exceeded
        """
        # Only handle batch creation
        if call_type != "acreate_batch":
            verbose_proxy_logger.debug(
                f"Batch rate limiter: Not handling batch creation rate limiting for call type: {call_type}"
            )
            return data

        verbose_proxy_logger.debug(
            "Batch rate limiter: Handling batch creation rate limiting"
        )

        try:
            # Extract input_file_id from data
            input_file_id = data.get("input_file_id")
            if not input_file_id:
                verbose_proxy_logger.debug(
                    "No input_file_id in batch request, skipping rate limiting"
                )
                return data

            # Get custom_llm_provider for token counting
            custom_llm_provider = data.get("custom_llm_provider", "openai")

            # Count tokens and requests from input file
            verbose_proxy_logger.debug(
                f"Counting tokens from batch input file: {input_file_id}"
            )
            batch_usage = await self.count_input_file_usage(
                file_id=input_file_id,
                custom_llm_provider=custom_llm_provider,
                user_api_key_dict=user_api_key_dict,
            )

            verbose_proxy_logger.debug(
                f"Batch input file usage - Tokens: {batch_usage.total_tokens}, "
                f"Requests: {batch_usage.request_count}"
            )

            # Store batch usage in data for later reference
            data["_batch_token_count"] = batch_usage.total_tokens
            data["_batch_request_count"] = batch_usage.request_count

            # Directly increment counters by batch amounts (check happens atomically)
            # This will raise HTTPException if limits are exceeded
            await self._check_and_increment_batch_counters(
                user_api_key_dict=user_api_key_dict,
                data=data,
                batch_usage=batch_usage,
            )

            verbose_proxy_logger.debug(
                "Batch rate limit check passed, counters incremented"
            )
            return data

        except HTTPException:
            # Re-raise HTTP exceptions (rate limit exceeded)
            raise
        except Exception as e:
            verbose_proxy_logger.error(
                f"Error in batch rate limiting: {str(e)}", exc_info=True
            )
            # Don't block the request if rate limiting fails
            return data