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

import httpx

import litellm
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.types.router import GenericLiteLLMParams
from litellm.types.vector_stores import (
    BaseVectorStoreAuthCredentials,
    VectorStoreCreateOptionalRequestParams,
    VectorStoreCreateResponse,
    VectorStoreIndexEndpoints,
    VectorStoreResultContent,
    VectorStoreSearchOptionalRequestParams,
    VectorStoreSearchResponse,
    VectorStoreSearchResult,
)

if TYPE_CHECKING:
    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj

    LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
    LiteLLMLoggingObj = Any


class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM):
    """
    Configuration for Azure AI Search Vector Store

    This implementation uses the Azure AI Search API for vector store operations.
    Supports vector search with embeddings generated via litellm.embeddings.
    """

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

    def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints:
        return {
            "read": [("GET", "/docs/search"), ("POST", "/docs/search")],
            "write": [("PUT", "/docs")],
        }

    def get_auth_credentials(
        self, litellm_params: dict
    ) -> BaseVectorStoreAuthCredentials:
        api_key = litellm_params.get("api_key")
        if api_key is None:
            raise ValueError("api_key is required")

        return {
            "headers": {
                "api-key": api_key,
            }
        }

    def validate_environment(
        self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]
    ) -> dict:
        basic_headers = self._base_validate_azure_environment(headers, litellm_params)
        basic_headers.update({"Content-Type": "application/json"})
        return basic_headers

    def get_complete_url(
        self,
        api_base: Optional[str],
        litellm_params: dict,
    ) -> str:
        """
        Get the base endpoint for Azure AI Search API

        Expected format: https://{search_service_name}.search.windows.net
        """
        if api_base:
            return api_base.rstrip("/")

        # Get search service name from litellm_params
        search_service_name = litellm_params.get("azure_search_service_name")

        if not search_service_name:
            raise ValueError(
                "Azure AI Search service name is required. "
                "Provide it via litellm_params['azure_search_service_name'] or api_base parameter"
            )

        # Azure AI Search endpoint
        return f"https://{search_service_name}.search.windows.net"

    def transform_search_vector_store_request(
        self,
        vector_store_id: str,
        query: Union[str, List[str]],
        vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
        api_base: str,
        litellm_logging_obj: LiteLLMLoggingObj,
        litellm_params: dict,
        extra_body: Optional[Dict[str, Any]] = None,
    ) -> Tuple[str, Dict[str, Any]]:
        """
        Transform search request for Azure AI Search API

        Generates embeddings using litellm.embeddings and constructs Azure AI Search request
        """
        # Convert query to string if it's a list
        if isinstance(query, list):
            query = " ".join(query)

        # Get embedding model from litellm_params (required)
        embedding_model = litellm_params.get("litellm_embedding_model")
        if not embedding_model:
            raise ValueError(
                "embedding_model is required in litellm_params for Azure AI Search. "
                "Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'"
            )

        embedding_config = litellm_params.get("litellm_embedding_config", {})
        if not embedding_config:
            raise ValueError(
                "embedding_config is required in litellm_params for Azure AI Search. "
                "Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}"
            )

        # Get vector field name (defaults to contentVector)
        vector_field = litellm_params.get("azure_search_vector_field", "contentVector")

        # Get top_k (number of results to return)
        top_k = vector_store_search_optional_params.get("top_k", 10)

        # Generate embedding for the query using litellm.embeddings
        try:
            embedding_response = litellm.embedding(
                model=embedding_model,
                input=[query],
                **embedding_config,
            )
            query_vector = embedding_response.data[0]["embedding"]
        except Exception as e:
            raise Exception(f"Failed to generate embedding for query: {str(e)}")

        # Azure AI Search endpoint for search
        index_name = vector_store_id  # vector_store_id is the index name
        url = f"{api_base}/indexes/{index_name}/docs/search?api-version=2024-07-01"

        # Build the request body for Azure AI Search with vector search
        request_body = {
            "search": "*",  # Get all documents (filtered by vector similarity)
            "vectorQueries": [
                {
                    "vector": query_vector,
                    "fields": vector_field,
                    "kind": "vector",
                    "k": top_k,  # Number of nearest neighbors to return
                }
            ],
            "select": "id,content",  # Fields to return (customize based on schema)
            "top": top_k,
        }

        #########################################################
        # Update logging object with details of the request
        #########################################################
        litellm_logging_obj.model_call_details["input"] = query
        litellm_logging_obj.model_call_details["embedding_model"] = embedding_model
        litellm_logging_obj.model_call_details["top_k"] = top_k

        return url, request_body

    def transform_search_vector_store_response(
        self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj
    ) -> VectorStoreSearchResponse:
        """
        Transform Azure AI Search API response to standard vector store search response

        Handles the format from Azure AI Search which returns:
        {
            "value": [
                {
                    "id": "...",
                    "content": "...",
                    "@search.score": 0.95,
                    ... (other fields)
                }
            ]
        }
        """
        try:
            response_json = response.json()

            # Extract results from Azure AI Search API response
            results = response_json.get("value", [])

            # Transform results to standard format
            search_results: List[VectorStoreSearchResult] = []
            for result in results:
                # Extract document ID
                document_id = result.get("id", "")

                # Extract text content
                text_content = result.get("content", "")

                content = [
                    VectorStoreResultContent(
                        text=text_content,
                        type="text",
                    )
                ]

                # Get the search score (relevance score from Azure AI Search)
                score = result.get("@search.score", 0.0)

                # Use document ID as both file_id and filename
                file_id = document_id
                filename = f"Document {document_id}"

                # Build attributes with all available metadata
                # Exclude system fields and already-processed fields
                attributes = {}
                for key, value in result.items():
                    if key not in ["id", "content", "contentVector", "@search.score"]:
                        attributes[key] = value

                # Always include document_id in attributes
                attributes["document_id"] = document_id

                result_obj = VectorStoreSearchResult(
                    score=score,
                    content=content,
                    file_id=file_id,
                    filename=filename,
                    attributes=attributes,
                )
                search_results.append(result_obj)

            return VectorStoreSearchResponse(
                object="vector_store.search_results.page",
                search_query=litellm_logging_obj.model_call_details.get("input", ""),
                data=search_results,
            )

        except Exception as e:
            raise self.get_error_class(
                error_message=str(e),
                status_code=response.status_code,
                headers=response.headers,
            )

    def transform_create_vector_store_request(
        self,
        vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams,
        api_base: str,
    ) -> Tuple[str, Dict]:
        raise NotImplementedError

    def transform_create_vector_store_response(
        self, response: httpx.Response
    ) -> VectorStoreCreateResponse:
        raise NotImplementedError
