"""
LiteLLM SDK Functions for Creating and Searching Vector Stores
"""

import asyncio
import contextvars
from functools import partial
from typing import Any, Coroutine, Dict, List, Optional, Union

import httpx

import litellm
from litellm.constants import request_timeout
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.router import GenericLiteLLMParams
from litellm.types.vector_stores import (
    VectorStoreCreateOptionalRequestParams,
    VectorStoreCreateResponse,
    VectorStoreFileCounts,
    VectorStoreResultContent,
    VectorStoreSearchOptionalRequestParams,
    VectorStoreSearchResponse,
    VectorStoreSearchResult,
)
from litellm.utils import ProviderConfigManager, client
from litellm.vector_stores.utils import VectorStoreRequestUtils

####### ENVIRONMENT VARIABLES ###################
# Initialize any necessary instances or variables here
base_llm_http_handler = BaseLLMHTTPHandler()
#################################################


def mock_vector_store_search_response(
    mock_results: Optional[List[VectorStoreSearchResult]] = None,
):
    """Mock response for vector store search"""
    if mock_results is None:
        mock_results = [
            VectorStoreSearchResult(
                score=0.95,
                content=[
                    VectorStoreResultContent(
                        text="This is a sample search result from the vector store.",
                        type="text",
                    )
                ],
            )
        ]

    return VectorStoreSearchResponse(
        object="vector_store.search_results.page",
        search_query="sample query",
        data=mock_results,
    )


def mock_vector_store_create_response(
    mock_response: Optional[VectorStoreCreateResponse] = None,
):
    """Mock response for vector store create"""
    if mock_response is None:
        mock_response = VectorStoreCreateResponse(
            id="vs_mock123",
            object="vector_store",
            created_at=1699061776,
            name="Mock Vector Store",
            bytes=0,
            file_counts=VectorStoreFileCounts(
                in_progress=0,
                completed=0,
                failed=0,
                cancelled=0,
                total=0,
            ),
            status="completed",
            expires_after=None,
            expires_at=None,
            last_active_at=None,
            metadata=None,
        )

    return mock_response


@client
async def acreate(
    name: Optional[str] = None,
    file_ids: Optional[List[str]] = None,
    expires_after: Optional[Dict] = None,
    chunking_strategy: Optional[Dict] = None,
    metadata: Optional[Dict[str, str]] = None,
    # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
    # The extra values given here take precedence over values defined on the client or passed to this method.
    extra_headers: Optional[Dict[str, Any]] = None,
    extra_query: Optional[Dict[str, Any]] = None,
    extra_body: Optional[Dict[str, Any]] = None,
    timeout: Optional[Union[float, httpx.Timeout]] = None,
    # LiteLLM specific params,
    custom_llm_provider: Optional[str] = None,
    **kwargs,
) -> VectorStoreCreateResponse:
    """
    Async: Create a vector store.
    """
    local_vars = locals()
    try:
        loop = asyncio.get_event_loop()
        kwargs["acreate"] = True

        # get custom llm provider so we can use this for mapping exceptions
        if custom_llm_provider is None:
            custom_llm_provider = "openai"  # Default to OpenAI for vector stores

        func = partial(
            create,
            name=name,
            file_ids=file_ids,
            expires_after=expires_after,
            chunking_strategy=chunking_strategy,
            metadata=metadata,
            extra_headers=extra_headers,
            extra_query=extra_query,
            extra_body=extra_body,
            timeout=timeout,
            custom_llm_provider=custom_llm_provider,
            **kwargs,
        )

        ctx = contextvars.copy_context()
        func_with_context = partial(ctx.run, func)
        init_response = await loop.run_in_executor(None, func_with_context)

        if asyncio.iscoroutine(init_response):
            response = await init_response
        else:
            response = init_response

        return response
    except Exception as e:
        raise litellm.exception_type(
            model=None,
            custom_llm_provider=custom_llm_provider,
            original_exception=e,
            completion_kwargs=local_vars,
            extra_kwargs=kwargs,
        )


@client
def create(
    name: Optional[str] = None,
    file_ids: Optional[List[str]] = None,
    expires_after: Optional[Dict] = None,
    chunking_strategy: Optional[Dict] = None,
    metadata: Optional[Dict[str, str]] = None,
    # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
    # The extra values given here take precedence over values defined on the client or passed to this method.
    extra_headers: Optional[Dict[str, Any]] = None,
    extra_query: Optional[Dict[str, Any]] = None,
    extra_body: Optional[Dict[str, Any]] = None,
    timeout: Optional[Union[float, httpx.Timeout]] = None,
    # LiteLLM specific params,
    custom_llm_provider: Optional[str] = None,
    **kwargs,
) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]:
    """
    Create a vector store.

    Args:
        name: The name of the vector store.
        file_ids: A list of File IDs that the vector store should use.
        expires_after: The expiration policy for the vector store.
        chunking_strategy: The chunking strategy used to chunk the file(s).
        metadata: Set of 16 key-value pairs that can be attached to an object.

    Returns:
        VectorStoreCreateResponse containing the created vector store details.
    """
    local_vars = locals()
    try:
        litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj")  # type: ignore
        litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
        _is_async = kwargs.pop("acreate", False) is True

        # get llm provider logic
        litellm_params = GenericLiteLLMParams(**kwargs)

        ## MOCK RESPONSE LOGIC
        if litellm_params.mock_response and isinstance(
            litellm_params.mock_response, dict
        ):
            return mock_vector_store_create_response(
                mock_response=VectorStoreCreateResponse(**litellm_params.mock_response)
            )

        # Default to OpenAI for vector stores
        if custom_llm_provider is None:
            custom_llm_provider = "openai"

        if "/" in custom_llm_provider:
            api_type, custom_llm_provider, _, _ = get_llm_provider(
                model=custom_llm_provider,
                custom_llm_provider=None,
                litellm_params=None,
            )
        else:
            api_type = None
            custom_llm_provider = custom_llm_provider

        # get provider config - using vector store custom logger for now
        vector_store_provider_config = (
            ProviderConfigManager.get_provider_vector_stores_config(
                provider=litellm.LlmProviders(custom_llm_provider),
                api_type=api_type,
            )
        )

        if vector_store_provider_config is None:
            raise ValueError(
                f"Vector store create is not supported for {custom_llm_provider}"
            )

        local_vars.update(kwargs)

        # Get VectorStoreCreateOptionalRequestParams with only valid parameters
        vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams = (
            VectorStoreRequestUtils.get_requested_vector_store_create_optional_param(
                local_vars
            )
        )

        # Pre Call logging
        litellm_logging_obj.update_environment_variables(
            model=None,
            optional_params={
                "name": name,
                **vector_store_create_optional_params,
            },
            litellm_params={
                "litellm_call_id": litellm_call_id,
            },
            custom_llm_provider=custom_llm_provider,
        )

        response = base_llm_http_handler.vector_store_create_handler(
            vector_store_create_optional_params=vector_store_create_optional_params,
            vector_store_provider_config=vector_store_provider_config,
            custom_llm_provider=custom_llm_provider,
            litellm_params=litellm_params,
            logging_obj=litellm_logging_obj,
            extra_headers=extra_headers,
            extra_body=extra_body,
            timeout=timeout or request_timeout,
            _is_async=_is_async,
            client=kwargs.get("client"),
        )

        return response
    except Exception as e:
        raise litellm.exception_type(
            model=None,
            custom_llm_provider=custom_llm_provider,
            original_exception=e,
            completion_kwargs=local_vars,
            extra_kwargs=kwargs,
        )


@client
async def asearch(
    vector_store_id: str,
    query: Union[str, List[str]],
    filters: Optional[Dict] = None,
    max_num_results: Optional[int] = None,
    ranking_options: Optional[Dict] = None,
    rewrite_query: Optional[bool] = None,
    # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
    # The extra values given here take precedence over values defined on the client or passed to this method.
    extra_headers: Optional[Dict[str, Any]] = None,
    extra_query: Optional[Dict[str, Any]] = None,
    extra_body: Optional[Dict[str, Any]] = None,
    timeout: Optional[Union[float, httpx.Timeout]] = None,
    # LiteLLM specific params,
    custom_llm_provider: Optional[str] = None,
    **kwargs,
) -> VectorStoreSearchResponse:
    """
    Async: Search a vector store for relevant chunks based on a query and file attributes filter.
    """
    local_vars = locals()

    try:
        loop = asyncio.get_event_loop()
        kwargs["asearch"] = True

        # get custom llm provider so we can use this for mapping exceptions
        if custom_llm_provider is None:
            custom_llm_provider = "openai"  # Default to OpenAI for vector stores

        func = partial(
            search,
            vector_store_id=vector_store_id,
            query=query,
            filters=filters,
            max_num_results=max_num_results,
            ranking_options=ranking_options,
            rewrite_query=rewrite_query,
            extra_headers=extra_headers,
            extra_query=extra_query,
            extra_body=extra_body,
            timeout=timeout,
            custom_llm_provider=custom_llm_provider,
            **kwargs,
        )

        ctx = contextvars.copy_context()
        func_with_context = partial(ctx.run, func)
        init_response = await loop.run_in_executor(None, func_with_context)

        if asyncio.iscoroutine(init_response):
            response = await init_response
        else:
            response = init_response

        return response
    except Exception as e:
        raise litellm.exception_type(
            model=None,
            custom_llm_provider=custom_llm_provider,
            original_exception=e,
            completion_kwargs=local_vars,
            extra_kwargs=kwargs,
        )


@client
def search(
    vector_store_id: str,
    query: Union[str, List[str]],
    filters: Optional[Dict] = None,
    max_num_results: Optional[int] = None,
    ranking_options: Optional[Dict] = None,
    rewrite_query: Optional[bool] = None,
    # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
    # The extra values given here take precedence over values defined on the client or passed to this method.
    extra_headers: Optional[Dict[str, Any]] = None,
    extra_query: Optional[Dict[str, Any]] = None,
    extra_body: Optional[Dict[str, Any]] = None,
    timeout: Optional[Union[float, httpx.Timeout]] = None,
    # LiteLLM specific params,
    custom_llm_provider: Optional[str] = None,
    **kwargs,
) -> Union[VectorStoreSearchResponse, Coroutine[Any, Any, VectorStoreSearchResponse]]:
    """
    Search a vector store for relevant chunks based on a query and file attributes filter.

    Args:
        vector_store_id: The ID of the vector store to search.
        query: A query string or array for the search.
        filters: Optional filter to apply based on file attributes.
        max_num_results: Maximum number of results to return (1-50, default 10).
        ranking_options: Optional ranking options for search.
        rewrite_query: Whether to rewrite the natural language query for vector search.

    Returns:
        VectorStoreSearchResponse containing the search results.
    """
    local_vars = locals()
    try:
        litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj")  # type: ignore
        litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
        _is_async = kwargs.pop("asearch", False) is True

        # pull credentials from registry if available
        vector_store_id_for_credentials = kwargs.get("vector_store_id", vector_store_id)
        if (
            litellm.vector_store_registry is not None
            and vector_store_id_for_credentials is not None
        ):
            try:
                registry_credentials = (
                    litellm.vector_store_registry.get_credentials_for_vector_store(
                        vector_store_id_for_credentials
                    )
                )
                kwargs.update(registry_credentials)
            except Exception:
                pass

        # get llm provider logic
        litellm_params = GenericLiteLLMParams(vector_store_id=vector_store_id, **kwargs)

        ## MOCK RESPONSE LOGIC
        if litellm_params.mock_response and isinstance(
            litellm_params.mock_response, (str, list)
        ):
            mock_results = None
            if isinstance(litellm_params.mock_response, list):
                mock_results = litellm_params.mock_response
            return mock_vector_store_search_response(mock_results=mock_results)

        # Default to OpenAI for vector stores
        if custom_llm_provider is None:
            custom_llm_provider = "openai"

        if "/" in custom_llm_provider:
            api_type, custom_llm_provider, _, _ = get_llm_provider(
                model=custom_llm_provider,
                custom_llm_provider=None,
                litellm_params=None,
            )
        else:
            api_type = None
            custom_llm_provider = custom_llm_provider

        # get provider config - using vector store custom logger for now
        vector_store_provider_config = (
            ProviderConfigManager.get_provider_vector_stores_config(
                provider=litellm.LlmProviders(custom_llm_provider),
                api_type=api_type,
            )
        )

        if vector_store_provider_config is None:
            raise ValueError(
                f"Vector store search is not supported for {custom_llm_provider}"
            )

        local_vars.update(kwargs)

        # Get VectorStoreSearchOptionalRequestParams with only valid parameters
        vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams = (
            VectorStoreRequestUtils.get_requested_vector_store_search_optional_param(
                local_vars,
                vector_store_provider_config=vector_store_provider_config,
            )
        )

        # Pre Call logging
        litellm_logging_obj.update_environment_variables(
            model=api_type,
            optional_params={
                "vector_store_id": vector_store_id,
                "query": query,
                **vector_store_search_optional_params,
            },
            litellm_params={
                "litellm_call_id": litellm_call_id,
                "vector_store_id": vector_store_id,
                **litellm_params.model_dump(exclude_none=True),
            },
            custom_llm_provider=custom_llm_provider,
        )

        response = base_llm_http_handler.vector_store_search_handler(
            vector_store_id=vector_store_id,
            query=query,
            vector_store_search_optional_params=vector_store_search_optional_params,
            vector_store_provider_config=vector_store_provider_config,
            custom_llm_provider=custom_llm_provider,
            litellm_params=litellm_params,
            logging_obj=litellm_logging_obj,
            extra_headers=extra_headers,
            extra_body=extra_body,
            timeout=timeout or request_timeout,
            _is_async=_is_async,
            client=kwargs.get("client"),
        )

        return response
    except Exception as e:
        raise litellm.exception_type(
            model=None,
            custom_llm_provider=custom_llm_provider,
            original_exception=e,
            completion_kwargs=local_vars,
            extra_kwargs=kwargs,
        )
