from typing import (
    Annotated,
    Any,  # noqa: TID251  # jsonify_object in proxy/utils.py is annotated with a bare dict
    Final,
    cast,  # noqa: TID251  # jsonify_object in proxy/utils.py is annotated with a bare dict
)

from fastapi import APIRouter, Depends, HTTPException, Request, Response

from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
    LiteLLM_ManagedVectorStore,
)
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.utils import jsonify_object
from litellm.proxy.vector_store_endpoints.management_endpoints import (
    _resolve_embedding_config,
)
from litellm.proxy.vector_store_endpoints.utils import (
    assert_proxy_admin_for_vector_store_index_management,
    assert_user_can_access_vector_store,
    get_litellm_managed_vector_store,
)
from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository
from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse
from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry

router: Final = APIRouter()
########################################################
# OpenAI Compatible Endpoints
########################################################


async def _update_request_data_with_litellm_managed_vector_store_registry(
    data: dict,
    vector_store_id: str,
    user_api_key_dict: UserAPIKeyAuth | None = None,
) -> dict:
    """
    Update the request data with the litellm managed vector store registry.

    Args:
        data: Request data to update
        vector_store_id: ID of the vector store
        user_api_key_dict: User API key authentication info for access control

    Raises:
        HTTPException: If user doesn't have access to the vector store
    """
    vector_store_to_run: Final[LiteLLM_ManagedVectorStore | None] = await get_litellm_managed_vector_store(
        vector_store_id=vector_store_id
    )
    if vector_store_to_run is not None:
        if user_api_key_dict is not None:
            await assert_user_can_access_vector_store(
                vector_store=vector_store_to_run,
                user_api_key_dict=user_api_key_dict,
            )

        if "custom_llm_provider" in vector_store_to_run:
            data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider")

        if "litellm_credential_name" in vector_store_to_run:
            data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name")

        if "litellm_params" in vector_store_to_run:
            litellm_params = vector_store_to_run.get("litellm_params", {}) or {}
            # Resolve ``litellm_embedding_config`` here, at request-handling
            # time, instead of at row-creation time. The resolved
            # ``api_key`` / ``api_base`` / ``api_version`` lives only in
            # this per-request ``data`` dict and is never persisted.
            # Legacy rows that already carry a resolved (cleartext)
            # ``litellm_embedding_config`` skip the lookup and pass through
            # unchanged so the embed call keeps working.
            embedding_model: Final = litellm_params.get("litellm_embedding_model")
            if embedding_model and not litellm_params.get("litellm_embedding_config"):
                from litellm.proxy.proxy_server import prisma_client

                resolved_config: Final = await _resolve_embedding_config(
                    embedding_model=embedding_model, prisma_client=prisma_client
                )
                if resolved_config:
                    # Build a fresh dict via spread instead of mutating
                    # ``litellm_params`` in place — the registry hands back
                    # a reference to its cached object, so an in-place
                    # update would persist the resolved cleartext into the
                    # in-memory cache for the lifetime of the process.
                    litellm_params = {
                        **litellm_params,
                        "litellm_embedding_config": resolved_config,
                    }
            data.update(litellm_params)
    return data


@router.post(
    "/v1/vector_stores/{vector_store_id:path}/search",
    dependencies=[Depends(user_api_key_auth)],
)
@router.post(
    "/vector_stores/{vector_store_id:path}/search",
    dependencies=[Depends(user_api_key_auth)],
)
async def vector_store_search(
    request: Request,
    vector_store_id: str,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Search a vector store.

    API Reference:
    https://platform.openai.com/docs/api-reference/vector-stores/search
    """
    from litellm.proxy.proxy_server import (
        _read_request_body,
        general_settings,
        llm_router,
        proxy_config,
        proxy_logging_obj,
        select_data_generator,
        user_api_base,
        user_max_tokens,
        user_model,
        user_request_timeout,
        user_temperature,
        version,
    )

    data = await _read_request_body(request=request)
    data["vector_store_id"] = vector_store_id

    # Check for legacy vector store registry (non-managed vector stores)
    data = await _update_request_data_with_litellm_managed_vector_store_registry(
        data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
    )

    # The managed_vector_stores pre-call hook will handle:
    # 1. Decoding managed vector store IDs
    # 2. Extracting model and provider resource ID
    # 3. Setting up proper routing
    # 4. Authentication checks

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="avector_store_search",
            proxy_logging_obj=proxy_logging_obj,
            llm_router=llm_router,
            general_settings=general_settings,
            proxy_config=proxy_config,
            select_data_generator=select_data_generator,
            model=None,
            user_model=user_model,
            user_temperature=user_temperature,
            user_request_timeout=user_request_timeout,
            user_max_tokens=user_max_tokens,
            user_api_base=user_api_base,
            version=version,
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=proxy_logging_obj,
            version=version,
        )


@router.post("/v1/vector_stores", dependencies=[Depends(user_api_key_auth)])
@router.post("/vector_stores", dependencies=[Depends(user_api_key_auth)])
async def vector_store_create(
    request: Request,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Create a vector store.

    API Reference:
    https://platform.openai.com/docs/api-reference/vector-stores/create

    Supports target_model_names parameter for creating vector stores across multiple models:
    ```json
    {
        "name": "my-vector-store",
        "target_model_names": "gpt-4,gemini-2.0"
    }
    ```
    """
    from litellm.proxy.proxy_server import (
        _read_request_body,
        general_settings,
        llm_router,
        proxy_config,
        proxy_logging_obj,
        select_data_generator,
        user_api_base,
        user_max_tokens,
        user_model,
        user_request_timeout,
        user_temperature,
        version,
    )

    data: Final = await _read_request_body(request=request)

    # Check for target_model_names parameter
    target_model_names: Final = data.pop("target_model_names", None)

    if target_model_names:
        # Use managed vector stores for multi-model support
        if isinstance(target_model_names, str):
            target_model_names_list = [m.strip() for m in target_model_names.split(",")]
        elif isinstance(target_model_names, list):
            target_model_names_list = target_model_names
        else:
            raise HTTPException(
                status_code=400,
                detail="target_model_names must be a comma-separated string or list of model names",
            )

        # Get managed vector stores hook
        managed_vector_stores: Final[Any] = proxy_logging_obj.get_proxy_hook("managed_vector_stores")
        if managed_vector_stores is None:
            raise HTTPException(
                status_code=500,
                detail="Managed vector stores not configured. Please ensure the proxy is initialized with database support.",
            )

        if llm_router is None:
            raise HTTPException(
                status_code=500,
                detail="LLM Router not initialized. Ensure models are added to proxy.",
            )

        # Create vector store across multiple models
        response: Final = await managed_vector_stores.acreate_vector_store(
            create_request=data,
            llm_router=llm_router,
            target_model_names_list=target_model_names_list,
            litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
            user_api_key_dict=user_api_key_dict,
        )

        return response

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="avector_store_create",
            proxy_logging_obj=proxy_logging_obj,
            llm_router=llm_router,
            general_settings=general_settings,
            proxy_config=proxy_config,
            select_data_generator=select_data_generator,
            model=None,
            user_model=user_model,
            user_temperature=user_temperature,
            user_request_timeout=user_request_timeout,
            user_max_tokens=user_max_tokens,
            user_api_base=user_api_base,
            version=version,
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=proxy_logging_obj,
            version=version,
        )


@router.get("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
@router.get("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
async def vector_store_retrieve(
    request: Request,
    vector_store_id: str,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Retrieve a vector store.

    API Reference:
    https://platform.openai.com/docs/api-reference/vector-stores/retrieve
    """
    from litellm.proxy.proxy_server import (
        general_settings,
        llm_router,
        proxy_config,
        proxy_logging_obj,
        select_data_generator,
        user_api_base,
        user_max_tokens,
        user_model,
        user_request_timeout,
        user_temperature,
        version,
    )

    data = {"vector_store_id": vector_store_id}

    data = await _update_request_data_with_litellm_managed_vector_store_registry(
        data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
    )

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="avector_store_retrieve",
            proxy_logging_obj=proxy_logging_obj,
            llm_router=llm_router,
            general_settings=general_settings,
            proxy_config=proxy_config,
            select_data_generator=select_data_generator,
            model=None,
            user_model=user_model,
            user_temperature=user_temperature,
            user_request_timeout=user_request_timeout,
            user_max_tokens=user_max_tokens,
            user_api_base=user_api_base,
            version=version,
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=proxy_logging_obj,
            version=version,
        )


@router.get("/v1/vector_stores", dependencies=[Depends(user_api_key_auth)])
@router.get("/vector_stores", dependencies=[Depends(user_api_key_auth)])
async def vector_store_list(
    request: Request,
    fastapi_response: Response,
    after: str | None = None,
    before: str | None = None,
    limit: int | None = 20,
    order: str | None = "desc",
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    List vector stores.

    API Reference:
    https://platform.openai.com/docs/api-reference/vector-stores/list
    """
    from litellm.proxy.proxy_server import (
        general_settings,
        llm_router,
        proxy_config,
        proxy_logging_obj,
        select_data_generator,
        user_api_base,
        user_max_tokens,
        user_model,
        user_request_timeout,
        user_temperature,
        version,
    )

    data: Final[dict] = {}
    if after is not None:
        data["after"] = after
    if before is not None:
        data["before"] = before
    if limit is not None:
        data["limit"] = limit
    if order is not None:
        data["order"] = order

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="avector_store_list",
            proxy_logging_obj=proxy_logging_obj,
            llm_router=llm_router,
            general_settings=general_settings,
            proxy_config=proxy_config,
            select_data_generator=select_data_generator,
            model=None,
            user_model=user_model,
            user_temperature=user_temperature,
            user_request_timeout=user_request_timeout,
            user_max_tokens=user_max_tokens,
            user_api_base=user_api_base,
            version=version,
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=proxy_logging_obj,
            version=version,
        )


@router.post("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
@router.post("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
async def vector_store_update(
    request: Request,
    vector_store_id: str,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Update a vector store.

    API Reference:
    https://platform.openai.com/docs/api-reference/vector-stores/modify
    """
    from litellm.proxy.proxy_server import (
        _read_request_body,
        general_settings,
        llm_router,
        proxy_config,
        proxy_logging_obj,
        select_data_generator,
        user_api_base,
        user_max_tokens,
        user_model,
        user_request_timeout,
        user_temperature,
        version,
    )

    data = await _read_request_body(request=request)
    if "vector_store_id" not in data:
        data["vector_store_id"] = vector_store_id

    data = await _update_request_data_with_litellm_managed_vector_store_registry(
        data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
    )

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="avector_store_update",
            proxy_logging_obj=proxy_logging_obj,
            llm_router=llm_router,
            general_settings=general_settings,
            proxy_config=proxy_config,
            select_data_generator=select_data_generator,
            model=None,
            user_model=user_model,
            user_temperature=user_temperature,
            user_request_timeout=user_request_timeout,
            user_max_tokens=user_max_tokens,
            user_api_base=user_api_base,
            version=version,
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=proxy_logging_obj,
            version=version,
        )


@router.delete("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
@router.delete("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
async def vector_store_delete(
    request: Request,
    vector_store_id: str,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Delete a vector store.

    API Reference:
    https://platform.openai.com/docs/api-reference/vector-stores/delete
    """
    from litellm.proxy.proxy_server import (
        general_settings,
        llm_router,
        proxy_config,
        proxy_logging_obj,
        select_data_generator,
        user_api_base,
        user_max_tokens,
        user_model,
        user_request_timeout,
        user_temperature,
        version,
    )

    data = {"vector_store_id": vector_store_id}

    data = await _update_request_data_with_litellm_managed_vector_store_registry(
        data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
    )

    processor: Final = ProxyBaseLLMRequestProcessing(data=data)
    try:
        return await processor.base_process_llm_request(
            request=request,
            fastapi_response=fastapi_response,
            user_api_key_dict=user_api_key_dict,
            route_type="avector_store_delete",
            proxy_logging_obj=proxy_logging_obj,
            llm_router=llm_router,
            general_settings=general_settings,
            proxy_config=proxy_config,
            select_data_generator=select_data_generator,
            model=None,
            user_model=user_model,
            user_temperature=user_temperature,
            user_request_timeout=user_request_timeout,
            user_max_tokens=user_max_tokens,
            user_api_base=user_api_base,
            version=version,
        )
    except Exception as e:
        raise await processor._handle_llm_api_exception(
            e=e,
            user_api_key_dict=user_api_key_dict,
            proxy_logging_obj=proxy_logging_obj,
            version=version,
        )


@router.post(
    "/v1/indexes",
    dependencies=[Depends(user_api_key_auth)],
)
async def index_create(
    request: Request,
    index_create_request: IndexCreateRequest,
    fastapi_response: Response,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Create an index. Just writes the index to the database.

    ```bash
    curl -L -X POST 'http://0.0.0.0:4000/v1/indexes' \
        -H 'Content-Type: application/json' \
        -H 'Authorization: Bearer sk-1234' \
        -d '{
            "index_name": "dall-e-3",
            "litellm_params": {
                "vector_store_index": "real-index-name",
                "vector_store_name": "azure-ai-search"
            }
        }'
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    assert_proxy_admin_for_vector_store_index_management(
        user_api_key_dict,
        operation="create",
    )

    if prisma_client is None:
        raise HTTPException(
            status_code=500,
            detail=CommonProxyErrors.db_not_connected_error.value,
        )
    ## 1. check if index already exists
    existing_index: Final = await ManagedVectorStoreIndexRepository(prisma_client).table.find_unique(
        where={"index_name": index_create_request.index_name}
    )

    ## 2. set created_by and updated_by

    if existing_index is not None:
        raise HTTPException(
            status_code=400,
            detail=f"Index {index_create_request.index_name} already exists",
        )

    ## 2. create index
    index_data: Final = index_create_request.model_dump(exclude_none=True)
    index_data["created_by"] = user_api_key_dict.user_id
    index_data["updated_by"] = user_api_key_dict.user_id
    new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create(
        data=cast(  # cast-ok: jsonify_object deep-copies a model_dump, so keys are str and values plain objects
            "dict[str, object]", jsonify_object(index_data)
        )
    )

    return new_index.model_dump()


@router.get(
    "/v1/indexes",
    dependencies=[Depends(user_api_key_auth)],
    response_model=IndexListResponse,
)
async def index_list(
    user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> IndexListResponse:
    """
    List all vector store indexes. Proxy admin only.

    ```bash
    curl -L -X GET 'http://0.0.0.0:4000/v1/indexes' \
        -H 'Authorization: Bearer sk-1234'
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    assert_proxy_admin_for_vector_store_index_management(
        user_api_key_dict,
        operation="list",
    )

    if prisma_client is None:
        raise HTTPException(
            status_code=500,
            detail=CommonProxyErrors.db_not_connected_error.value,
        )

    indexes: Final = await VectorStoreIndexRegistry._get_vector_store_indexes_from_db(prisma_client)
    return IndexListResponse(data=indexes)