"""
Transformation utilities for Vertex AI RAG Engine.

Handles transforming LiteLLM's unified formats to Vertex AI RAG Engine API format.
"""

from typing import Any, Dict, Optional

from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.rag import RAGChunkingStrategy


class VertexAIRAGTransformation(VertexBase):
    """
    Transformation class for Vertex AI RAG Engine API.

    Handles:
    - Converting unified chunking_strategy to Vertex AI format
    - Building import request payloads
    - Transforming responses
    """

    def __init__(self):
        super().__init__()

    def get_import_rag_files_url(
        self,
        vertex_project: str,
        vertex_location: str,
        corpus_id: str,
    ) -> str:
        """
        Get the URL for importing RAG files.

        Note: The REST endpoint for importRagFiles may not be publicly available.
        Vertex AI RAG Engine primarily uses gRPC-based SDK.
        """
        base_url = get_vertex_base_url(vertex_location)
        return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{corpus_id}:importRagFiles"

    def get_retrieve_contexts_url(
        self,
        vertex_project: str,
        vertex_location: str,
    ) -> str:
        """Get the URL for retrieving contexts (search)."""
        base_url = get_vertex_base_url(vertex_location)
        return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}:retrieveContexts"

    def transform_chunking_strategy_to_vertex_format(
        self,
        chunking_strategy: Optional[RAGChunkingStrategy],
    ) -> Dict[str, Any]:
        """
        Transform LiteLLM's unified chunking_strategy to Vertex AI RAG format.

        LiteLLM format (RAGChunkingStrategy):
            {
                "chunk_size": 1000,
                "chunk_overlap": 200,
                "separators": ["\n\n", "\n", " ", ""]
            }

        Vertex AI RAG format (TransformationConfig):
            {
                "chunking_config": {
                    "chunk_size": 1000,
                    "chunk_overlap": 200
                }
            }

        Note: Vertex AI doesn't support custom separators in the same way,
        so we only transform chunk_size and chunk_overlap.
        """
        if not chunking_strategy:
            return {
                "chunking_config": {
                    "chunk_size": DEFAULT_CHUNK_SIZE,
                    "chunk_overlap": DEFAULT_CHUNK_OVERLAP,
                }
            }

        chunk_size = chunking_strategy.get("chunk_size", DEFAULT_CHUNK_SIZE)
        chunk_overlap = chunking_strategy.get("chunk_overlap", DEFAULT_CHUNK_OVERLAP)

        # Log if separators are provided (not supported by Vertex AI)
        if chunking_strategy.get("separators"):
            verbose_logger.warning(
                "Vertex AI RAG Engine does not support custom separators. "
                "The 'separators' parameter will be ignored."
            )

        return {
            "chunking_config": {
                "chunk_size": chunk_size,
                "chunk_overlap": chunk_overlap,
            }
        }

    def build_import_rag_files_request(
        self,
        gcs_uri: str,
        chunking_strategy: Optional[RAGChunkingStrategy] = None,
    ) -> Dict[str, Any]:
        """
        Build the request payload for importing RAG files.

        Args:
            gcs_uri: GCS URI of the file to import (e.g., gs://bucket/path/file.txt)
            chunking_strategy: LiteLLM unified chunking config

        Returns:
            Request payload dict for importRagFiles API
        """
        transformation_config = self.transform_chunking_strategy_to_vertex_format(
            chunking_strategy
        )

        return {
            "import_rag_files_config": {
                "gcs_source": {
                    "uris": [gcs_uri]
                },
                "rag_file_transformation_config": transformation_config,
            }
        }

    def get_auth_headers(
        self,
        vertex_credentials: Optional[str] = None,
        vertex_project: Optional[str] = None,
    ) -> Dict[str, str]:
        """
        Get authentication headers for Vertex AI API calls.

        Uses the base class method to get credentials.
        """
        credentials = self.get_vertex_ai_credentials(
            {"vertex_credentials": vertex_credentials}
        )
        project = vertex_project or self.get_vertex_ai_project({})

        access_token, _ = self._ensure_access_token(
            credentials=credentials,
            project_id=project,
            custom_llm_provider="vertex_ai",
        )

        return {
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json",
        }