from typing import (
    Any,
    Callable,
    cast,
    Dict,
    Sequence,
    Optional,
    Type,
    TypeVar,
    Tuple,
)
import fastapi
import orjson
from anyio import (
    to_thread,
    CapacityLimiter,
)
from fastapi import FastAPI as _FastAPI, Response, Request, Body
from fastapi.responses import JSONResponse, ORJSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.routing import APIRoute
from fastapi import HTTPException, status

from chromadb.api.configuration import CollectionConfigurationInternal
from pydantic import BaseModel
from chromadb.api.types import (
    Embedding,
    GetResult,
    QueryResult,
    Embeddings,
    convert_list_embeddings_to_np,
)
from chromadb.auth import UserIdentity
from chromadb.auth import (
    AuthzAction,
    AuthzResource,
    ServerAuthenticationProvider,
    ServerAuthorizationProvider,
)
from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings, System
from chromadb.api import ServerAPI
from chromadb.errors import (
    ChromaError,
    InvalidDimensionException,
    InvalidHTTPVersion,
    RateLimitError,
    QuotaError,
)
from chromadb.server import Server
from chromadb.server.fastapi.types import (
    AddEmbedding,
    CreateDatabase,
    CreateTenant,
    DeleteEmbedding,
    GetEmbedding,
    QueryEmbedding,
    CreateCollection,
    UpdateCollection,
    UpdateEmbedding,
)
from starlette.datastructures import Headers
import logging

from chromadb.telemetry.product.events import ServerStartEvent
from chromadb.utils.fastapi import fastapi_json_response, string_to_uuid as _uuid
from opentelemetry import trace

from chromadb.telemetry.opentelemetry.fastapi import instrument_fastapi
from chromadb.types import Database, Tenant
from chromadb.telemetry.product import ServerContext, ProductTelemetryClient
from chromadb.telemetry.opentelemetry import (
    OpenTelemetryClient,
    OpenTelemetryGranularity,
    add_attributes_to_current_span,
    trace_method,
)
from chromadb.types import Collection as CollectionModel

logger = logging.getLogger(__name__)


def use_route_names_as_operation_ids(app: _FastAPI) -> None:
    """
    Simplify operation IDs so that generated API clients have simpler function
    names.
    Should be called only after all routes have been added.
    """
    for route in app.routes:
        if isinstance(route, APIRoute):
            route.operation_id = route.name


async def add_trace_id_to_response_middleware(
    request: Request, call_next: Callable[[Request], Any]
) -> Response:
    trace_id = trace.get_current_span().get_span_context().trace_id
    response = await call_next(request)
    response.headers["Chroma-Trace-Id"] = format(trace_id, "x")
    return response


async def catch_exceptions_middleware(
    request: Request, call_next: Callable[[Request], Any]
) -> Response:
    try:
        return await call_next(request)
    except ChromaError as e:
        return fastapi_json_response(e)
    except ValueError as e:
        return JSONResponse(
            content={"error": "InvalidArgumentError", "message": str(e)},
            status_code=400,
        )
    except TypeError as e:
        return JSONResponse(
            content={"error": "InvalidArgumentError", "message": str(e)},
            status_code=400,
        )
    except Exception as e:
        logger.exception(e)
        return JSONResponse(content={"error": repr(e)}, status_code=500)


async def check_http_version_middleware(
    request: Request, call_next: Callable[[Request], Any]
) -> Response:
    http_version = request.scope.get("http_version")
    if http_version not in ["1.1", "2"]:
        raise InvalidHTTPVersion(f"HTTP version {http_version} is not supported")
    return await call_next(request)


D = TypeVar("D", bound=BaseModel, contravariant=True)


def validate_model(model: Type[D], data: Any) -> D:  # type: ignore
    """Used for backward compatibility with Pydantic 1.x"""
    try:
        return model.model_validate(data)  # pydantic 2.x
    except AttributeError:
        return model.parse_obj(data)  # pydantic 1.x


class ChromaAPIRouter(fastapi.APIRouter):  # type: ignore
    # A simple subclass of fastapi's APIRouter which treats URLs with a
    # trailing "/" the same as URLs without. Docs will only contain URLs
    # without trailing "/"s.
    def add_api_route(self, path: str, *args: Any, **kwargs: Any) -> None:
        # If kwargs["include_in_schema"] isn't passed OR is True, we should
        # only include the non-"/" path. If kwargs["include_in_schema"] is
        # False, include neither.
        exclude_from_schema = (
            "include_in_schema" in kwargs and not kwargs["include_in_schema"]
        )

        def include_in_schema(path: str) -> bool:
            nonlocal exclude_from_schema
            return not exclude_from_schema and not path.endswith("/")

        kwargs["include_in_schema"] = include_in_schema(path)
        super().add_api_route(path, *args, **kwargs)

        if path.endswith("/"):
            path = path[:-1]
        else:
            path = path + "/"

        kwargs["include_in_schema"] = include_in_schema(path)
        super().add_api_route(path, *args, **kwargs)


class FastAPI(Server):
    def __init__(self, settings: Settings):
        ProductTelemetryClient.SERVER_CONTEXT = ServerContext.FASTAPI
        # https://fastapi.tiangolo.com/advanced/custom-response/#use-orjsonresponse
        self._app = fastapi.FastAPI(debug=True, default_response_class=ORJSONResponse)
        self._system = System(settings)
        self._api: ServerAPI = self._system.instance(ServerAPI)
        self._opentelemetry_client = self._api.require(OpenTelemetryClient)
        self._capacity_limiter = CapacityLimiter(
            settings.chroma_server_thread_pool_size
        )
        self._system.start()

        self._app.middleware("http")(check_http_version_middleware)
        self._app.middleware("http")(catch_exceptions_middleware)
        self._app.middleware("http")(add_trace_id_to_response_middleware)
        self._app.add_middleware(
            CORSMiddleware,
            allow_headers=["*"],
            allow_origins=settings.chroma_server_cors_allow_origins,
            allow_methods=["*"],
        )
        self._app.add_exception_handler(QuotaError, self.quota_exception_handler)
        self._app.add_exception_handler(
            RateLimitError, self.rate_limit_exception_handler
        )

        self._app.on_event("shutdown")(self.shutdown)

        self.authn_provider = None
        if settings.chroma_server_authn_provider:
            self.authn_provider = self._system.require(ServerAuthenticationProvider)

        self.authz_provider = None
        if settings.chroma_server_authz_provider:
            self.authz_provider = self._system.require(ServerAuthorizationProvider)

        self.router = ChromaAPIRouter()

        self.setup_v1_routes()
        self.setup_v2_routes()

        self._app.include_router(self.router)

        use_route_names_as_operation_ids(self._app)
        instrument_fastapi(self._app)
        telemetry_client = self._system.instance(ProductTelemetryClient)
        telemetry_client.capture(ServerStartEvent())

    def setup_v2_routes(self) -> None:
        self.router.add_api_route("/api/v2", self.root, methods=["GET"])
        self.router.add_api_route("/api/v2/reset", self.reset, methods=["POST"])
        self.router.add_api_route("/api/v2/version", self.version, methods=["GET"])
        self.router.add_api_route("/api/v2/heartbeat", self.heartbeat, methods=["GET"])
        self.router.add_api_route(
            "/api/v2/pre-flight-checks", self.pre_flight_checks, methods=["GET"]
        )

        self.router.add_api_route(
            "/api/v2/auth/identity",
            self.get_user_identity,
            methods=["GET"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases",
            self.create_database,
            methods=["POST"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}",
            self.get_database,
            methods=["GET"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v2/tenants",
            self.create_tenant,
            methods=["POST"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v2/tenants/{tenant}",
            self.get_tenant,
            methods=["GET"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections",
            self.list_collections,
            methods=["GET"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections_count",
            self.count_collections,
            methods=["GET"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections",
            self.create_collection,
            methods=["POST"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_id}/add",
            self.add,
            methods=["POST"],
            status_code=status.HTTP_201_CREATED,
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_id}/update",
            self.update,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_id}/upsert",
            self.upsert,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_id}/get",
            self.get,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_id}/delete",
            self.delete,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_id}/count",
            self.count,
            methods=["GET"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_id}/query",
            self.get_nearest_neighbors,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_name}",
            self.get_collection,
            methods=["GET"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_id}",
            self.update_collection,
            methods=["PUT"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v2/tenants/{tenant}/databases/{database_name}/collections/{collection_name}",
            self.delete_collection,
            methods=["DELETE"],
            response_model=None,
        )

    def shutdown(self) -> None:
        self._system.stop()

    def app(self) -> fastapi.FastAPI:
        return self._app

    async def rate_limit_exception_handler(
        self, request: Request, exc: RateLimitError
    ) -> JSONResponse:
        return JSONResponse(
            status_code=429,
            content={"message": "Rate limit exceeded."},
        )

    def root(self) -> Dict[str, int]:
        return {"nanosecond heartbeat": self._api.heartbeat()}

    async def quota_exception_handler(
        self, request: Request, exc: QuotaError
    ) -> JSONResponse:
        return JSONResponse(
            status_code=429,
            content={"message": exc.message()},
        )

    async def heartbeat(self) -> Dict[str, int]:
        return self.root()

    async def version(self) -> str:
        return self._api.get_version()

    @trace_method(
        "auth_request",
        OpenTelemetryGranularity.OPERATION,
    )
    def auth_request(
        self,
        headers: Headers,
        action: AuthzAction,
        tenant: Optional[str],
        database: Optional[str],
        collection: Optional[str],
    ) -> None:
        """
        Authenticates and authorizes the request based on the given headers
        and other parameters. If the request cannot be authenticated or cannot
        be authorized (with the configured providers), raises an HTTP 401.
        """
        if not self.authn_provider:
            add_attributes_to_current_span(
                {
                    "tenant": tenant,
                    "database": database,
                    "collection": collection,
                }
            )
            return

        user_identity = self.authn_provider.authenticate_or_raise(dict(headers))

        if not self.authz_provider:
            return

        authz_resource = AuthzResource(
            tenant=tenant,
            database=database,
            collection=collection,
        )

        self.authz_provider.authorize_or_raise(user_identity, action, authz_resource)
        add_attributes_to_current_span(
            {
                "tenant": tenant,
                "database": database,
                "collection": collection,
            }
        )
        return

    @trace_method("FastAPI.get_user_identity", OpenTelemetryGranularity.OPERATION)
    async def get_user_identity(
        self,
        request: Request,
    ) -> UserIdentity:
        if not self.authn_provider:
            return UserIdentity(
                user_id="", tenant=DEFAULT_TENANT, databases=[DEFAULT_DATABASE]
            )

        return cast(
            UserIdentity,
            await to_thread.run_sync(
                lambda: cast(ServerAuthenticationProvider, self.authn_provider).authenticate_or_raise(dict(request.headers))  # type: ignore
            ),
        )

    @trace_method("FastAPI.create_database", OpenTelemetryGranularity.OPERATION)
    async def create_database(
        self,
        request: Request,
        tenant: str,
        body: CreateDatabase = Body(...),
    ) -> None:
        def process_create_database(
            tenant: str, headers: Headers, raw_body: bytes
        ) -> None:
            db = validate_model(CreateDatabase, orjson.loads(raw_body))

            self.auth_request(
                headers,
                AuthzAction.CREATE_DATABASE,
                tenant,
                db.name,
                None,
            )

            return self._api.create_database(db.name, tenant)

        await to_thread.run_sync(
            process_create_database,
            tenant,
            request.headers,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.get_database", OpenTelemetryGranularity.OPERATION)
    async def get_database(
        self,
        request: Request,
        database_name: str,
        tenant: str,
    ) -> Database:
        self.auth_request(
            request.headers,
            AuthzAction.GET_DATABASE,
            tenant,
            database_name,
            None,
        )

        return cast(
            Database,
            await to_thread.run_sync(
                self._api.get_database,
                database_name,
                tenant,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.create_tenant", OpenTelemetryGranularity.OPERATION)
    async def create_tenant(
        self, request: Request, body: CreateTenant = Body(...)
    ) -> None:
        def process_create_tenant(request: Request, raw_body: bytes) -> None:
            tenant = validate_model(CreateTenant, orjson.loads(raw_body))

            self.auth_request(
                request.headers,
                AuthzAction.CREATE_TENANT,
                tenant.name,
                None,
                None,
            )

            return self._api.create_tenant(tenant.name)

        await to_thread.run_sync(
            process_create_tenant,
            request,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.get_tenant", OpenTelemetryGranularity.OPERATION)
    async def get_tenant(
        self,
        request: Request,
        tenant: str,
    ) -> Tenant:
        self.auth_request(
            request.headers,
            AuthzAction.GET_TENANT,
            tenant,
            None,
            None,
        )

        return cast(
            Tenant,
            await to_thread.run_sync(
                self._api.get_tenant,
                tenant,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.list_collections", OpenTelemetryGranularity.OPERATION)
    async def list_collections(
        self,
        request: Request,
        tenant: str,
        database_name: str,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
    ) -> Sequence[CollectionModel]:
        self.auth_request(
            request.headers,
            AuthzAction.LIST_COLLECTIONS,
            tenant,
            database_name,
            None,
        )

        add_attributes_to_current_span({"tenant": tenant})

        api_collection_models = cast(
            Sequence[CollectionModel],
            await to_thread.run_sync(
                self._api.list_collections,
                limit,
                offset,
                tenant,
                database_name,
                limiter=self._capacity_limiter,
            ),
        )

        return api_collection_models

    @trace_method("FastAPI.count_collections", OpenTelemetryGranularity.OPERATION)
    async def count_collections(
        self,
        request: Request,
        tenant: str,
        database_name: str,
    ) -> int:
        self.auth_request(
            request.headers,
            AuthzAction.COUNT_COLLECTIONS,
            tenant,
            database_name,
            None,
        )

        add_attributes_to_current_span({"tenant": tenant})

        return cast(
            int,
            await to_thread.run_sync(
                self._api.count_collections,
                tenant,
                database_name,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.create_collection", OpenTelemetryGranularity.OPERATION)
    async def create_collection(
        self,
        request: Request,
        tenant: str,
        database_name: str,
        body: CreateCollection = Body(...),
    ) -> CollectionModel:
        def process_create_collection(
            request: Request, tenant: str, database: str, raw_body: bytes
        ) -> CollectionModel:
            create = validate_model(CreateCollection, orjson.loads(raw_body))
            configuration = (
                CollectionConfigurationInternal()
                if not create.configuration
                else CollectionConfigurationInternal.from_json(create.configuration)
            )

            self.auth_request(
                request.headers,
                AuthzAction.CREATE_COLLECTION,
                tenant,
                database,
                create.name,
            )

            add_attributes_to_current_span({"tenant": tenant})

            return self._api.create_collection(
                name=create.name,
                configuration=configuration,
                metadata=create.metadata,
                get_or_create=create.get_or_create,
                tenant=tenant,
                database=database,
            )

        api_collection_model = cast(
            CollectionModel,
            await to_thread.run_sync(
                process_create_collection,
                request,
                tenant,
                database_name,
                await request.body(),
                limiter=self._capacity_limiter,
            ),
        )
        return api_collection_model

    @trace_method("FastAPI.get_collection", OpenTelemetryGranularity.OPERATION)
    async def get_collection(
        self,
        request: Request,
        tenant: str,
        database_name: str,
        collection_name: str,
    ) -> CollectionModel:
        self.auth_request(
            request.headers,
            AuthzAction.GET_COLLECTION,
            tenant,
            database_name,
            collection_name,
        )

        add_attributes_to_current_span({"tenant": tenant})

        api_collection_model = cast(
            CollectionModel,
            await to_thread.run_sync(
                self._api.get_collection,
                collection_name,
                tenant,
                database_name,
                limiter=self._capacity_limiter,
            ),
        )
        return api_collection_model

    @trace_method("FastAPI.update_collection", OpenTelemetryGranularity.OPERATION)
    async def update_collection(
        self,
        tenant: str,
        database_name: str,
        collection_id: str,
        request: Request,
        body: UpdateCollection = Body(...),
    ) -> None:
        def process_update_collection(
            request: Request, collection_id: str, raw_body: bytes
        ) -> None:
            update = validate_model(UpdateCollection, orjson.loads(raw_body))
            self.auth_request(
                request.headers,
                AuthzAction.UPDATE_COLLECTION,
                tenant,
                database_name,
                collection_id,
            )
            add_attributes_to_current_span({"tenant": tenant})
            return self._api._modify(
                id=_uuid(collection_id),
                new_name=update.new_name,
                new_metadata=update.new_metadata,
                tenant=tenant,
                database=database_name,
            )

        await to_thread.run_sync(
            process_update_collection,
            request,
            collection_id,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.delete_collection", OpenTelemetryGranularity.OPERATION)
    async def delete_collection(
        self,
        request: Request,
        collection_name: str,
        tenant: str,
        database_name: str,
    ) -> None:
        self.auth_request(
            request.headers,
            AuthzAction.DELETE_COLLECTION,
            tenant,
            database_name,
            collection_name,
        )
        add_attributes_to_current_span({"tenant": tenant})

        await to_thread.run_sync(
            self._api.delete_collection,
            collection_name,
            tenant,
            database_name,
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.add", OpenTelemetryGranularity.OPERATION)
    async def add(
        self,
        request: Request,
        tenant: str,
        database_name: str,
        collection_id: str,
        body: AddEmbedding = Body(...),
    ) -> bool:
        try:

            def process_add(request: Request, raw_body: bytes) -> bool:
                add = validate_model(AddEmbedding, orjson.loads(raw_body))
                self.auth_request(
                    request.headers,
                    AuthzAction.ADD,
                    tenant,
                    database_name,
                    collection_id,
                )
                add_attributes_to_current_span({"tenant": tenant})
                return self._api._add(
                    collection_id=_uuid(collection_id),
                    ids=add.ids,
                    embeddings=cast(
                        Embeddings,
                        convert_list_embeddings_to_np(add.embeddings)
                        if add.embeddings
                        else None,
                    ),
                    metadatas=add.metadatas,  # type: ignore
                    documents=add.documents,  # type: ignore
                    uris=add.uris,  # type: ignore
                    tenant=tenant,
                    database=database_name,
                )

            return cast(
                bool,
                await to_thread.run_sync(
                    process_add,
                    request,
                    await request.body(),
                    limiter=self._capacity_limiter,
                ),
            )
        except InvalidDimensionException as e:
            raise HTTPException(status_code=500, detail=str(e))

    @trace_method("FastAPI.update", OpenTelemetryGranularity.OPERATION)
    async def update(
        self,
        request: Request,
        tenant: str,
        database_name: str,
        collection_id: str,
        body: UpdateEmbedding = Body(...),
    ) -> None:
        def process_update(request: Request, raw_body: bytes) -> bool:
            update = validate_model(UpdateEmbedding, orjson.loads(raw_body))

            self.auth_request(
                request.headers,
                AuthzAction.UPDATE,
                tenant,
                database_name,
                collection_id,
            )
            add_attributes_to_current_span({"tenant": tenant})

            return self._api._update(
                collection_id=_uuid(collection_id),
                ids=update.ids,
                embeddings=convert_list_embeddings_to_np(update.embeddings)
                if update.embeddings
                else None,
                metadatas=update.metadatas,  # type: ignore
                documents=update.documents,  # type: ignore
                uris=update.uris,  # type: ignore
                tenant=tenant,
                database=database_name,
            )

        await to_thread.run_sync(
            process_update,
            request,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.upsert", OpenTelemetryGranularity.OPERATION)
    async def upsert(
        self,
        request: Request,
        tenant: str,
        database_name: str,
        collection_id: str,
        body: AddEmbedding = Body(...),
    ) -> None:
        def process_upsert(request: Request, raw_body: bytes) -> bool:
            upsert = validate_model(AddEmbedding, orjson.loads(raw_body))

            self.auth_request(
                request.headers,
                AuthzAction.UPSERT,
                tenant,
                database_name,
                collection_id,
            )
            add_attributes_to_current_span({"tenant": tenant})

            return self._api._upsert(
                collection_id=_uuid(collection_id),
                ids=upsert.ids,
                embeddings=cast(
                    Embeddings,
                    convert_list_embeddings_to_np(upsert.embeddings)
                    if upsert.embeddings
                    else None,
                ),
                metadatas=upsert.metadatas,  # type: ignore
                documents=upsert.documents,  # type: ignore
                uris=upsert.uris,  # type: ignore
                tenant=tenant,
                database=database_name,
            )

        await to_thread.run_sync(
            process_upsert,
            request,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.get", OpenTelemetryGranularity.OPERATION)
    async def get(
        self,
        collection_id: str,
        tenant: str,
        database_name: str,
        request: Request,
        body: GetEmbedding = Body(...),
    ) -> GetResult:
        def process_get(request: Request, raw_body: bytes) -> GetResult:
            get = validate_model(GetEmbedding, orjson.loads(raw_body))
            self.auth_request(
                request.headers,
                AuthzAction.GET,
                tenant,
                database_name,
                collection_id,
            )
            add_attributes_to_current_span({"tenant": tenant})
            return self._api._get(
                collection_id=_uuid(collection_id),
                ids=get.ids,
                where=get.where,
                sort=get.sort,
                limit=get.limit,
                offset=get.offset,
                where_document=get.where_document,
                include=get.include,
                tenant=tenant,
                database=database_name,
            )

        get_result = cast(
            GetResult,
            await to_thread.run_sync(
                process_get,
                request,
                await request.body(),
                limiter=self._capacity_limiter,
            ),
        )

        if get_result["embeddings"] is not None:
            get_result["embeddings"] = [
                cast(Embedding, embedding).tolist()
                for embedding in get_result["embeddings"]
            ]

        return get_result

    @trace_method("FastAPI.delete", OpenTelemetryGranularity.OPERATION)
    async def delete(
        self,
        collection_id: str,
        tenant: str,
        database_name: str,
        request: Request,
        body: DeleteEmbedding = Body(...),
    ) -> None:
        def process_delete(request: Request, raw_body: bytes) -> None:
            delete = validate_model(DeleteEmbedding, orjson.loads(raw_body))
            self.auth_request(
                request.headers,
                AuthzAction.DELETE,
                tenant,
                database_name,
                collection_id,
            )
            add_attributes_to_current_span({"tenant": tenant})
            return self._api._delete(
                collection_id=_uuid(collection_id),
                ids=delete.ids,
                where=delete.where,
                where_document=delete.where_document,
                tenant=tenant,
                database=database_name,
            )

        await to_thread.run_sync(
            process_delete,
            request,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.count", OpenTelemetryGranularity.OPERATION)
    async def count(
        self,
        request: Request,
        tenant: str,
        database_name: str,
        collection_id: str,
    ) -> int:
        self.auth_request(
            request.headers,
            AuthzAction.COUNT,
            tenant,
            database_name,
            collection_id,
        )
        add_attributes_to_current_span({"tenant": tenant})

        return cast(
            int,
            await to_thread.run_sync(
                self._api._count,
                _uuid(collection_id),
                tenant,
                database_name,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.reset", OpenTelemetryGranularity.OPERATION)
    async def reset(
        self,
        request: Request,
    ) -> bool:
        self.auth_request(
            request.headers,
            AuthzAction.RESET,
            None,
            None,
            None,
        )

        return cast(
            bool,
            await to_thread.run_sync(
                self._api.reset,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.get_nearest_neighbors", OpenTelemetryGranularity.OPERATION)
    async def get_nearest_neighbors(
        self,
        tenant: str,
        database_name: str,
        collection_id: str,
        request: Request,
        body: QueryEmbedding = Body(...),
    ) -> QueryResult:
        def process_query(request: Request, raw_body: bytes) -> QueryResult:
            query = validate_model(QueryEmbedding, orjson.loads(raw_body))

            self.auth_request(
                request.headers,
                AuthzAction.QUERY,
                tenant,
                database_name,
                collection_id,
            )
            add_attributes_to_current_span({"tenant": tenant})

            return self._api._query(
                collection_id=_uuid(collection_id),
                query_embeddings=cast(
                    Embeddings,
                    convert_list_embeddings_to_np(query.query_embeddings)
                    if query.query_embeddings
                    else None,
                ),
                n_results=query.n_results,
                where=query.where,
                where_document=query.where_document,
                include=query.include,
                tenant=tenant,
                database=database_name,
            )

        nnresult = cast(
            QueryResult,
            await to_thread.run_sync(
                process_query,
                request,
                await request.body(),
                limiter=self._capacity_limiter,
            ),
        )

        if nnresult["embeddings"] is not None:
            nnresult["embeddings"] = [
                [cast(Embedding, embedding).tolist() for embedding in result]
                for result in nnresult["embeddings"]
            ]

        return nnresult

    async def pre_flight_checks(self) -> Dict[str, Any]:
        def process_pre_flight_checks() -> Dict[str, Any]:
            return {
                "max_batch_size": self._api.get_max_batch_size(),
            }

        return cast(
            Dict[str, Any],
            await to_thread.run_sync(
                process_pre_flight_checks,
                limiter=self._capacity_limiter,
            ),
        )

    # =========================================================================
    # OLD ROUTES FOR BACKWARDS COMPATIBILITY — WILL BE REMOVED
    # =========================================================================
    def setup_v1_routes(self) -> None:
        self.router.add_api_route("/api/v1", self.root, methods=["GET"])
        self.router.add_api_route("/api/v1/reset", self.reset, methods=["POST"])
        self.router.add_api_route("/api/v1/version", self.version, methods=["GET"])
        self.router.add_api_route("/api/v1/heartbeat", self.heartbeat, methods=["GET"])
        self.router.add_api_route(
            "/api/v1/pre-flight-checks", self.pre_flight_checks, methods=["GET"]
        )

        self.router.add_api_route(
            "/api/v1/databases",
            self.create_database_v1,
            methods=["POST"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v1/databases/{database}",
            self.get_database_v1,
            methods=["GET"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v1/tenants",
            self.create_tenant_v1,
            methods=["POST"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v1/tenants/{tenant}",
            self.get_tenant_v1,
            methods=["GET"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v1/collections",
            self.list_collections_v1,
            methods=["GET"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/count_collections",
            self.count_collections_v1,
            methods=["GET"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections",
            self.create_collection_v1,
            methods=["POST"],
            response_model=None,
        )

        self.router.add_api_route(
            "/api/v1/collections/{collection_id}/add",
            self.add_v1,
            methods=["POST"],
            status_code=status.HTTP_201_CREATED,
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_id}/update",
            self.update_v1,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_id}/upsert",
            self.upsert_v1,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_id}/get",
            self.get_v1,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_id}/delete",
            self.delete_v1,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_id}/count",
            self.count_v1,
            methods=["GET"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_id}/query",
            self.get_nearest_neighbors_v1,
            methods=["POST"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_name}",
            self.get_collection_v1,
            methods=["GET"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_id}",
            self.update_collection_v1,
            methods=["PUT"],
            response_model=None,
        )
        self.router.add_api_route(
            "/api/v1/collections/{collection_name}",
            self.delete_collection_v1,
            methods=["DELETE"],
            response_model=None,
        )

    @trace_method(
        "auth_and_get_tenant_and_database_for_request_v1",
        OpenTelemetryGranularity.OPERATION,
    )
    def auth_and_get_tenant_and_database_for_request(
        self,
        headers: Headers,
        action: AuthzAction,
        tenant: Optional[str],
        database: Optional[str],
        collection: Optional[str],
    ) -> Tuple[Optional[str], Optional[str]]:
        """
        Authenticates and authorizes the request based on the given headers
        and other parameters. If the request cannot be authenticated or cannot
        be authorized (with the configured providers), raises an HTTP 401.

        If the request is authenticated and authorized, returns the tenant and
        database to be used for the request. These will differ from the passed
        tenant and database if and only if:
        - The request is authenticated
        - chroma_overwrite_singleton_tenant_database_access_from_auth = True
        - The passed tenant or database are None or default_{tenant, database}
            (can be overwritten separately)
        - The user has access to a single tenant and/or single database.
        """
        if not self.authn_provider:
            add_attributes_to_current_span(
                {
                    "tenant": tenant,
                    "database": database,
                    "collection": collection,
                }
            )
            return (tenant, database)

        user_identity = self.authn_provider.authenticate_or_raise(dict(headers))

        (
            new_tenant,
            new_database,
        ) = self.authn_provider.singleton_tenant_database_if_applicable(user_identity)

        if (not tenant or tenant == DEFAULT_TENANT) and new_tenant:
            tenant = new_tenant
        if (not database or database == DEFAULT_DATABASE) and new_database:
            database = new_database

        if not self.authz_provider:
            return (tenant, database)

        authz_resource = AuthzResource(
            tenant=tenant,
            database=database,
            collection=collection,
        )

        self.authz_provider.authorize_or_raise(user_identity, action, authz_resource)
        add_attributes_to_current_span(
            {
                "tenant": tenant,
                "database": database,
                "collection": collection,
            }
        )
        return (tenant, database)

    @trace_method("FastAPI.create_database_v1", OpenTelemetryGranularity.OPERATION)
    async def create_database_v1(
        self,
        request: Request,
        tenant: str = DEFAULT_TENANT,
        body: CreateDatabase = Body(...),
    ) -> None:
        def process_create_database(
            tenant: str, headers: Headers, raw_body: bytes
        ) -> None:
            db = validate_model(CreateDatabase, orjson.loads(raw_body))

            (
                maybe_tenant,
                maybe_database,
            ) = self.auth_and_get_tenant_and_database_for_request(
                headers,
                AuthzAction.CREATE_DATABASE,
                tenant,
                db.name,
                None,
            )
            if maybe_tenant:
                tenant = maybe_tenant
            if maybe_database:
                db.name = maybe_database

            return self._api.create_database(db.name, tenant)

        await to_thread.run_sync(
            process_create_database,
            tenant,
            request.headers,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.get_database_v1", OpenTelemetryGranularity.OPERATION)
    async def get_database_v1(
        self,
        request: Request,
        database: str,
        tenant: str = DEFAULT_TENANT,
    ) -> Database:
        (
            maybe_tenant,
            maybe_database,
        ) = self.auth_and_get_tenant_and_database_for_request(
            request.headers,
            AuthzAction.GET_DATABASE,
            tenant,
            database,
            None,
        )
        if maybe_tenant:
            tenant = maybe_tenant
        if maybe_database:
            database = maybe_database

        return cast(
            Database,
            await to_thread.run_sync(
                self._api.get_database,
                database,
                tenant,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.create_tenant_v1", OpenTelemetryGranularity.OPERATION)
    async def create_tenant_v1(
        self, request: Request, body: CreateTenant = Body(...)
    ) -> None:
        def process_create_tenant(request: Request, raw_body: bytes) -> None:
            tenant = validate_model(CreateTenant, orjson.loads(raw_body))

            maybe_tenant, _ = self.auth_and_get_tenant_and_database_for_request(
                request.headers,
                AuthzAction.CREATE_TENANT,
                tenant.name,
                None,
                None,
            )
            if maybe_tenant:
                tenant.name = maybe_tenant

            return self._api.create_tenant(tenant.name)

        await to_thread.run_sync(
            process_create_tenant,
            request,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.get_tenant_v1", OpenTelemetryGranularity.OPERATION)
    async def get_tenant_v1(
        self,
        request: Request,
        tenant: str,
    ) -> Tenant:
        maybe_tenant, _ = self.auth_and_get_tenant_and_database_for_request(
            request.headers,
            AuthzAction.GET_TENANT,
            tenant,
            None,
            None,
        )
        if maybe_tenant:
            tenant = maybe_tenant

        return cast(
            Tenant,
            await to_thread.run_sync(
                self._api.get_tenant,
                tenant,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.list_collections_v1", OpenTelemetryGranularity.OPERATION)
    async def list_collections_v1(
        self,
        request: Request,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        tenant: str = DEFAULT_TENANT,
        database: str = DEFAULT_DATABASE,
    ) -> Sequence[CollectionModel]:
        (
            maybe_tenant,
            maybe_database,
        ) = self.auth_and_get_tenant_and_database_for_request(
            request.headers,
            AuthzAction.LIST_COLLECTIONS,
            tenant,
            database,
            None,
        )
        if maybe_tenant:
            tenant = maybe_tenant
        if maybe_database:
            database = maybe_database

        api_collection_models = cast(
            Sequence[CollectionModel],
            await to_thread.run_sync(
                self._api.list_collections,
                limit,
                offset,
                tenant,
                database,
                limiter=self._capacity_limiter,
            ),
        )

        return api_collection_models

    @trace_method("FastAPI.count_collections_v1", OpenTelemetryGranularity.OPERATION)
    async def count_collections_v1(
        self,
        request: Request,
        tenant: str = DEFAULT_TENANT,
        database: str = DEFAULT_DATABASE,
    ) -> int:
        (
            maybe_tenant,
            maybe_database,
        ) = self.auth_and_get_tenant_and_database_for_request(
            request.headers,
            AuthzAction.COUNT_COLLECTIONS,
            tenant,
            database,
            None,
        )
        if maybe_tenant:
            tenant = maybe_tenant
        if maybe_database:
            database = maybe_database

        return cast(
            int,
            await to_thread.run_sync(
                self._api.count_collections,
                tenant,
                database,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.create_collection_v1", OpenTelemetryGranularity.OPERATION)
    async def create_collection_v1(
        self,
        request: Request,
        tenant: str = DEFAULT_TENANT,
        database: str = DEFAULT_DATABASE,
        body: CreateCollection = Body(...),
    ) -> CollectionModel:
        def process_create_collection(
            request: Request, tenant: str, database: str, raw_body: bytes
        ) -> CollectionModel:
            create = validate_model(CreateCollection, orjson.loads(raw_body))
            configuration = (
                CollectionConfigurationInternal()
                if not create.configuration
                else CollectionConfigurationInternal.from_json(create.configuration)
            )

            (
                maybe_tenant,
                maybe_database,
            ) = self.auth_and_get_tenant_and_database_for_request(
                request.headers,
                AuthzAction.CREATE_COLLECTION,
                tenant,
                database,
                create.name,
            )
            if maybe_tenant:
                tenant = maybe_tenant
            if maybe_database:
                database = maybe_database

            return self._api.create_collection(
                name=create.name,
                configuration=configuration,
                metadata=create.metadata,
                get_or_create=create.get_or_create,
                tenant=tenant,
                database=database,
            )

        api_collection_model = cast(
            CollectionModel,
            await to_thread.run_sync(
                process_create_collection,
                request,
                tenant,
                database,
                await request.body(),
                limiter=self._capacity_limiter,
            ),
        )
        return api_collection_model

    @trace_method("FastAPI.get_collection_v1", OpenTelemetryGranularity.OPERATION)
    async def get_collection_v1(
        self,
        request: Request,
        collection_name: str,
        tenant: str = DEFAULT_TENANT,
        database: str = DEFAULT_DATABASE,
    ) -> CollectionModel:
        (
            maybe_tenant,
            maybe_database,
        ) = self.auth_and_get_tenant_and_database_for_request(
            request.headers,
            AuthzAction.GET_COLLECTION,
            tenant,
            database,
            collection_name,
        )
        if maybe_tenant:
            tenant = maybe_tenant
        if maybe_database:
            database = maybe_database

        api_collection_model = cast(
            CollectionModel,
            await to_thread.run_sync(
                self._api.get_collection,
                collection_name,
                tenant,
                database,
                limiter=self._capacity_limiter,
            ),
        )
        return api_collection_model

    @trace_method("FastAPI.update_collection_v1", OpenTelemetryGranularity.OPERATION)
    async def update_collection_v1(
        self, collection_id: str, request: Request, body: UpdateCollection = Body(...)
    ) -> None:
        def process_update_collection(
            request: Request, collection_id: str, raw_body: bytes
        ) -> None:
            update = validate_model(UpdateCollection, orjson.loads(raw_body))
            self.auth_and_get_tenant_and_database_for_request(
                request.headers,
                AuthzAction.UPDATE_COLLECTION,
                None,
                None,
                collection_id,
            )
            return self._api._modify(
                id=_uuid(collection_id),
                new_name=update.new_name,
                new_metadata=update.new_metadata,
            )

        await to_thread.run_sync(
            process_update_collection,
            request,
            collection_id,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.delete_collection_v1", OpenTelemetryGranularity.OPERATION)
    async def delete_collection_v1(
        self,
        request: Request,
        collection_name: str,
        tenant: str = DEFAULT_TENANT,
        database: str = DEFAULT_DATABASE,
    ) -> None:
        (
            maybe_tenant,
            maybe_database,
        ) = self.auth_and_get_tenant_and_database_for_request(
            request.headers,
            AuthzAction.DELETE_COLLECTION,
            tenant,
            database,
            collection_name,
        )
        if maybe_tenant:
            tenant = maybe_tenant
        if maybe_database:
            database = maybe_database

        await to_thread.run_sync(
            self._api.delete_collection,
            collection_name,
            tenant,
            database,
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.add_v1", OpenTelemetryGranularity.OPERATION)
    async def add_v1(
        self, request: Request, collection_id: str, body: AddEmbedding = Body(...)
    ) -> bool:
        try:

            def process_add(request: Request, raw_body: bytes) -> bool:
                add = validate_model(AddEmbedding, orjson.loads(raw_body))
                self.auth_and_get_tenant_and_database_for_request(
                    request.headers,
                    AuthzAction.ADD,
                    None,
                    None,
                    collection_id,
                )
                return self._api._add(
                    collection_id=_uuid(collection_id),
                    ids=add.ids,
                    embeddings=cast(
                        Embeddings,
                        convert_list_embeddings_to_np(add.embeddings)
                        if add.embeddings
                        else None,
                    ),
                    metadatas=add.metadatas,  # type: ignore
                    documents=add.documents,  # type: ignore
                    uris=add.uris,  # type: ignore
                )

            return cast(
                bool,
                await to_thread.run_sync(
                    process_add,
                    request,
                    await request.body(),
                    limiter=self._capacity_limiter,
                ),
            )
        except InvalidDimensionException as e:
            raise HTTPException(status_code=500, detail=str(e))

    @trace_method("FastAPI.update_v1", OpenTelemetryGranularity.OPERATION)
    async def update_v1(
        self, request: Request, collection_id: str, body: UpdateEmbedding = Body(...)
    ) -> None:
        def process_update(request: Request, raw_body: bytes) -> bool:
            update = validate_model(UpdateEmbedding, orjson.loads(raw_body))

            self.auth_and_get_tenant_and_database_for_request(
                request.headers,
                AuthzAction.UPDATE,
                None,
                None,
                collection_id,
            )

            return self._api._update(
                collection_id=_uuid(collection_id),
                ids=update.ids,
                embeddings=convert_list_embeddings_to_np(update.embeddings)
                if update.embeddings
                else None,
                metadatas=update.metadatas,  # type: ignore
                documents=update.documents,  # type: ignore
                uris=update.uris,  # type: ignore
            )

        await to_thread.run_sync(
            process_update,
            request,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.upsert_v1", OpenTelemetryGranularity.OPERATION)
    async def upsert_v1(
        self, request: Request, collection_id: str, body: AddEmbedding = Body(...)
    ) -> None:
        def process_upsert(request: Request, raw_body: bytes) -> bool:
            upsert = validate_model(AddEmbedding, orjson.loads(raw_body))

            self.auth_and_get_tenant_and_database_for_request(
                request.headers,
                AuthzAction.UPSERT,
                None,
                None,
                collection_id,
            )

            return self._api._upsert(
                collection_id=_uuid(collection_id),
                ids=upsert.ids,
                embeddings=cast(
                    Embeddings,
                    convert_list_embeddings_to_np(upsert.embeddings)
                    if upsert.embeddings
                    else None,
                ),
                metadatas=upsert.metadatas,  # type: ignore
                documents=upsert.documents,  # type: ignore
                uris=upsert.uris,  # type: ignore
            )

        await to_thread.run_sync(
            process_upsert,
            request,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.get_v1", OpenTelemetryGranularity.OPERATION)
    async def get_v1(
        self, collection_id: str, request: Request, body: GetEmbedding = Body(...)
    ) -> GetResult:
        def process_get(request: Request, raw_body: bytes) -> GetResult:
            get = validate_model(GetEmbedding, orjson.loads(raw_body))
            self.auth_and_get_tenant_and_database_for_request(
                request.headers,
                AuthzAction.GET,
                None,
                None,
                collection_id,
            )
            return self._api._get(
                collection_id=_uuid(collection_id),
                ids=get.ids,
                where=get.where,
                sort=get.sort,
                limit=get.limit,
                offset=get.offset,
                where_document=get.where_document,
                include=get.include,
            )

        get_result = cast(
            GetResult,
            await to_thread.run_sync(
                process_get,
                request,
                await request.body(),
                limiter=self._capacity_limiter,
            ),
        )

        if get_result["embeddings"] is not None:
            get_result["embeddings"] = [
                cast(Embedding, embedding).tolist()
                for embedding in get_result["embeddings"]
            ]

        return get_result

    @trace_method("FastAPI.delete_v1", OpenTelemetryGranularity.OPERATION)
    async def delete_v1(
        self, collection_id: str, request: Request, body: DeleteEmbedding = Body(...)
    ) -> None:
        def process_delete(request: Request, raw_body: bytes) -> None:
            delete = validate_model(DeleteEmbedding, orjson.loads(raw_body))
            self.auth_and_get_tenant_and_database_for_request(
                request.headers,
                AuthzAction.DELETE,
                None,
                None,
                collection_id,
            )
            return self._api._delete(
                collection_id=_uuid(collection_id),
                ids=delete.ids,
                where=delete.where,
                where_document=delete.where_document,
            )

        await to_thread.run_sync(
            process_delete,
            request,
            await request.body(),
            limiter=self._capacity_limiter,
        )

    @trace_method("FastAPI.count_v1", OpenTelemetryGranularity.OPERATION)
    async def count_v1(
        self,
        request: Request,
        collection_id: str,
    ) -> int:
        self.auth_and_get_tenant_and_database_for_request(
            request.headers,
            AuthzAction.COUNT,
            None,
            None,
            collection_id,
        )

        return cast(
            int,
            await to_thread.run_sync(
                self._api._count,
                _uuid(collection_id),
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method("FastAPI.reset_v1", OpenTelemetryGranularity.OPERATION)
    async def reset_v1(
        self,
        request: Request,
    ) -> bool:
        self.auth_and_get_tenant_and_database_for_request(
            request.headers,
            AuthzAction.RESET,
            None,
            None,
            None,
        )

        return cast(
            bool,
            await to_thread.run_sync(
                self._api.reset,
                limiter=self._capacity_limiter,
            ),
        )

    @trace_method(
        "FastAPI.get_nearest_neighbors_v1", OpenTelemetryGranularity.OPERATION
    )
    async def get_nearest_neighbors_v1(
        self,
        collection_id: str,
        request: Request,
        body: QueryEmbedding = Body(...),
    ) -> QueryResult:
        def process_query(request: Request, raw_body: bytes) -> QueryResult:
            query = validate_model(QueryEmbedding, orjson.loads(raw_body))

            self.auth_and_get_tenant_and_database_for_request(
                request.headers,
                AuthzAction.QUERY,
                None,
                None,
                collection_id,
            )

            return self._api._query(
                collection_id=_uuid(collection_id),
                query_embeddings=cast(
                    Embeddings,
                    convert_list_embeddings_to_np(query.query_embeddings)
                    if query.query_embeddings
                    else None,
                ),
                n_results=query.n_results,
                where=query.where,
                where_document=query.where_document,
                include=query.include,
            )

        nnresult = cast(
            QueryResult,
            await to_thread.run_sync(
                process_query,
                request,
                await request.body(),
                limiter=self._capacity_limiter,
            ),
        )

        if nnresult["embeddings"] is not None:
            nnresult["embeddings"] = [
                [cast(Embedding, embedding).tolist() for embedding in result]
                for result in nnresult["embeddings"]
            ]

        return nnresult

    # =========================================================================
