# Copyright DataStax, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
import json
import logging
import queue
import threading
import weakref
from collections.abc import AsyncGenerator, AsyncIterator
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from types import TracebackType
from typing import Any, Callable, Iterator, List, Union, cast

import httpx

from astrapy.core.api import APIRequestError, api_request, async_api_request
from astrapy.core.core_types import (
    API_DOC,
    API_RESPONSE,
    AsyncPaginableRequestMethod,
    PaginableRequestMethod,
)
from astrapy.core.defaults import (
    DEFAULT_AUTH_HEADER,
    DEFAULT_INSERT_NUM_DOCUMENTS,
    DEFAULT_JSON_API_PATH,
    DEFAULT_JSON_API_VERSION,
    DEFAULT_KEYSPACE_NAME,
)
from astrapy.core.utils import (
    TimeoutInfoWideType,
    convert_vector_to_floats,
    http_methods,
    make_payload,
    normalize_for_api,
    restore_from_api,
    to_httpx_timeout,
)

logger = logging.getLogger(__name__)


class _PrefetchIterator(Iterator[API_DOC]):
    def __init__(
        self,
        prefetched: int,
        request_method: PaginableRequestMethod,
        options: dict[str, Any] | None,
        raw_response_callback: Callable[[dict[str, Any]], None] | None = None,
    ):
        self.queue: queue.Queue[API_DOC | None] = queue.Queue(prefetched)
        self.request_method = request_method
        self.options = options
        self.raw_response_callback = raw_response_callback
        self.initialised = threading.Event()
        self.stop = threading.Event()
        self.thread = threading.Thread(
            target=_PrefetchIterator.queued_paginate, args=(weakref.proxy(self),)
        )
        self.thread.start()
        # wait until the exception handler in queued_paginate can deal with the
        # object being deleted.
        self.initialised.wait()

    def __iter__(self) -> Iterator[API_DOC]:
        return self

    @staticmethod
    def queue_put(
        q: queue.Queue[API_DOC | None],
        item: API_DOC | None,
        stop: threading.Event,
    ) -> None:
        while not stop.is_set():
            try:
                q.put(item, timeout=1)
                break
            except queue.Full:
                # Wait until there is space in the queue or the thread
                # is stopped
                pass

    def queued_paginate(self) -> None:
        self.initialised.set()
        try:
            for row in AstraDBCollection.paginate(
                request_method=self.request_method,
                options=self.options,
                raw_response_callback=self.raw_response_callback,
            ):
                self.queue_put(self.queue, row, self.stop)
        except ReferenceError:
            logger.debug("queued_paginate terminated")
            return
        finally:
            self.queue_put(self.queue, None, self.stop)
        logger.debug("queued_paginate end")

    def __next__(self) -> API_DOC:
        doc = self.queue.get()
        if doc is None:
            raise StopIteration
        return doc

    def __del__(self) -> None:
        self.stop.set()


class AstraDBCollection:
    # Initialize the shared httpx client as a class attribute
    client = httpx.Client()

    def __init__(
        self,
        collection_name: str,
        astra_db: AstraDB | None = None,
        token: str | None = None,
        api_endpoint: str | None = None,
        namespace: str | None = None,
        caller_name: str | None = None,
        caller_version: str | None = None,
        additional_headers: dict[str, str] = {},
    ) -> None:
        """
        Initialize an AstraDBCollection instance.

        Args:
            collection_name (str): The name of the collection.
            astra_db (AstraDB, optional): An instance of Astra DB.
            token (str, optional): Authentication token for Astra DB.
            api_endpoint (str, optional): API endpoint URL.
            namespace (str, optional): Namespace for the database.
            caller_name (str, optional): identity of the caller ("my_framework")
                If passing a client, its caller is used as fallback
            caller_version (str, optional): version of the caller code ("1.0.3")
                If passing a client, its caller is used as fallback
            additional_headers (Dict[str, str]): any further set of headers,
                in the form of key-value pairs, to be passed with the HTTP
                requests by this collection instance.
        """
        # Check for presence of the Astra DB object
        if astra_db is None:
            if api_endpoint is None:
                raise AssertionError("Must provide api_endpoint")

            astra_db = AstraDB(
                token=token,
                api_endpoint=api_endpoint,
                namespace=namespace,
                caller_name=caller_name,
                caller_version=caller_version,
            )
        else:
            # if astra_db passed, copy and apply possible overrides
            astra_db = astra_db.copy(
                token=token,
                api_endpoint=api_endpoint,
                namespace=namespace,
                caller_name=caller_name,
                caller_version=caller_version,
            )

        # Set the remaining instance attributes
        self.astra_db = astra_db
        self.caller_name: str | None = self.astra_db.caller_name
        self.caller_version: str | None = self.astra_db.caller_version
        self.additional_headers = additional_headers
        self.collection_name = collection_name
        self.base_path: str = f"{self.astra_db.base_path}/{self.collection_name}"

    def __repr__(self) -> str:
        return f'AstraDBCollection[astra_db="{self.astra_db}", collection_name="{self.collection_name}"]'

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, AstraDBCollection):
            return all(
                [
                    self.collection_name == other.collection_name,
                    self.astra_db == other.astra_db,
                    self.caller_name == other.caller_name,
                    self.caller_version == other.caller_version,
                    self.additional_headers == other.additional_headers,
                ]
            )
        else:
            return False

    def copy(
        self,
        *,
        collection_name: str | None = None,
        token: str | None = None,
        api_endpoint: str | None = None,
        api_path: str | None = None,
        api_version: str | None = None,
        namespace: str | None = None,
        caller_name: str | None = None,
        caller_version: str | None = None,
        additional_headers: dict[str, str] | None = None,
    ) -> AstraDBCollection:
        return AstraDBCollection(
            collection_name=collection_name or self.collection_name,
            astra_db=self.astra_db.copy(
                token=token,
                api_endpoint=api_endpoint,
                api_path=api_path,
                api_version=api_version,
                namespace=namespace,
                caller_name=caller_name,
                caller_version=caller_version,
            ),
            caller_name=caller_name or self.caller_name,
            caller_version=caller_version or self.caller_version,
            additional_headers=additional_headers or self.additional_headers,
        )

    def to_async(self) -> AsyncAstraDBCollection:
        return AsyncAstraDBCollection(
            collection_name=self.collection_name,
            astra_db=self.astra_db.to_async(),
            caller_name=self.caller_name,
            caller_version=self.caller_version,
            additional_headers=self.additional_headers,
        )

    def set_caller(
        self,
        caller_name: str | None = None,
        caller_version: str | None = None,
    ) -> None:
        self.astra_db.set_caller(
            caller_name=caller_name,
            caller_version=caller_version,
        )
        self.caller_name = caller_name
        self.caller_version = caller_version

    def _request(
        self,
        method: str = http_methods.POST,
        path: str | None = None,
        json_data: dict[str, Any] | None = None,
        url_params: dict[str, Any] | None = None,
        skip_error_check: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        direct_response = api_request(
            client=self.client,
            base_url=self.astra_db.base_url,
            auth_header=DEFAULT_AUTH_HEADER,
            token=self.astra_db.token,
            method=method,
            json_data=normalize_for_api(json_data),
            url_params=url_params,
            path=path,
            skip_error_check=skip_error_check,
            caller_name=self.caller_name,
            caller_version=self.caller_version,
            timeout=to_httpx_timeout(timeout_info),
            additional_headers=self.additional_headers,
        )
        response = restore_from_api(direct_response)
        return response

    def post_raw_request(
        self, body: dict[str, Any], timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        return self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=body,
            timeout_info=timeout_info,
        )

    def _get(
        self,
        path: str | None = None,
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE | None:
        full_path = f"{self.base_path}/{path}" if path else self.base_path
        response = self._request(
            method=http_methods.GET,
            path=full_path,
            url_params=options,
            timeout_info=timeout_info,
        )
        if isinstance(response, dict):
            return response
        return None

    def _put(
        self,
        path: str | None = None,
        document: API_RESPONSE | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        full_path = f"{self.base_path}/{path}" if path else self.base_path
        response = self._request(
            method=http_methods.PUT,
            path=full_path,
            json_data=document,
            timeout_info=timeout_info,
        )
        return response

    def _post(
        self,
        path: str | None = None,
        document: API_DOC | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        full_path = f"{self.base_path}/{path}" if path else self.base_path
        response = self._request(
            method=http_methods.POST,
            path=full_path,
            json_data=document,
            timeout_info=timeout_info,
        )
        return response

    def _recast_as_sort_projection(
        self, vector: list[float], fields: list[str] | None = None
    ) -> tuple[dict[str, Any], dict[str, Any] | None]:
        """
        Given a vector and optionally a list of fields,
        reformulate them as a sort, projection pair for regular
        'find'-like API calls (with basic validation as well).
        """
        # Must pass a vector
        if not vector:
            raise ValueError("Must pass a vector")

        # Edge case for field selection
        if fields and "$similarity" in fields:
            raise ValueError("Please use the `include_similarity` parameter")

        # Build the new vector parameter
        sort: dict[str, Any] = {"$vector": vector}

        # Build the new fields parameter
        # Note: do not leave projection={}, make it None
        # (or it will devour $similarity away in the API response)
        if fields is not None and len(fields) > 0:
            projection = {f: 1 for f in fields}
        else:
            projection = None

        return sort, projection

    def get(
        self, path: str | None = None, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE | None:
        """
        Retrieve a document from the collection by its path.

        Args:
            path (str, optional): The path of the document to retrieve.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The retrieved document.
        """
        return self._get(path=path, timeout_info=timeout_info)

    def find(
        self,
        filter: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        sort: dict[str, Any] | None = None,
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find documents in the collection that match the given filter.

        Args:
            filter (dict, optional): Criteria to filter documents.
            projection (dict, optional): Specifies the fields to return.
            sort (dict, optional): Specifies the order in which to return matching documents.
            options (dict, optional): Additional options for the query.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The query response containing matched documents.
        """
        json_query = make_payload(
            top_level="find",
            filter=filter,
            projection=projection,
            options=options,
            sort=sort,
        )

        response = self._post(document=json_query, timeout_info=timeout_info)

        return response

    def vector_find(
        self,
        vector: list[float],
        *,
        limit: int,
        filter: dict[str, Any] | None = None,
        fields: list[str] | None = None,
        include_similarity: bool = True,
        timeout_info: TimeoutInfoWideType = None,
    ) -> list[API_DOC]:
        """
        Perform a vector-based search in the collection.

        Args:
            vector (list): The vector to search with.
            limit (int): The maximum number of documents to return.
            filter (dict, optional): Criteria to filter documents.
            fields (list, optional): Specifies the fields to return.
            include_similarity (bool, optional): Whether to include similarity score in the result.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            list: A list of documents matching the vector search criteria.
        """
        # Must pass a limit
        if not limit:
            raise ValueError("Must pass a limit")

        # Pre-process the included arguments
        sort, projection = self._recast_as_sort_projection(
            convert_vector_to_floats(vector),
            fields=fields,
        )

        # Call the underlying find() method to search
        raw_find_result = self.find(
            filter=filter,
            projection=projection,
            sort=sort,
            options={
                "limit": limit,
                "includeSimilarity": include_similarity,
            },
            timeout_info=timeout_info,
        )

        return cast(List[API_DOC], raw_find_result["data"]["documents"])

    @staticmethod
    def paginate(
        *,
        request_method: PaginableRequestMethod,
        options: dict[str, Any] | None,
        raw_response_callback: Callable[[dict[str, Any]], None] | None = None,
    ) -> Iterator[API_DOC]:
        """
        Generate paginated results for a given database query method.

        Args:
            request_method (function): The database query method to paginate.
            options (dict, optional): Options for the database query.
            prefetched (int, optional): Number of pre-fetched documents.
            raw_response_callback: an optional callback invoked at each new
                response coming from the API. The only argument is the raw
                API response and the callback must return None.

        Yields:
            dict: The next document in the paginated result set.
        """
        _options = options or {}
        response0 = request_method(options=_options)
        if raw_response_callback:
            raw_response_callback(response0)
        next_page_state = response0["data"]["nextPageState"]
        options0 = _options

        for document in response0["data"]["documents"]:
            yield document
        while next_page_state is not None:
            options1 = {**options0, **{"pageState": next_page_state}}
            response1 = request_method(options=options1)
            if raw_response_callback:
                raw_response_callback(response1)
            for document in response1["data"]["documents"]:
                yield document
            next_page_state = response1["data"]["nextPageState"]

    def paginated_find(
        self,
        filter: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        sort: dict[str, Any] | None = None,
        options: dict[str, Any] | None = None,
        prefetched: int | None = None,
        timeout_info: TimeoutInfoWideType = None,
        raw_response_callback: Callable[[dict[str, Any]], None] | None = None,
    ) -> Iterator[API_DOC]:
        """
        Perform a paginated search in the collection.

        Args:
            filter (dict, optional): Criteria to filter documents.
            projection (dict, optional): Specifies the fields to return.
            sort (dict, optional): Specifies the order in which to return matching documents.
            options (dict, optional): Additional options for the query.
            prefetched (int, optional): Number of pre-fetched documents.
            timeout_info: a float, or a TimeoutInfo dict, for each
                single HTTP request.
                This is a paginated method, that issues several requests as it
                needs more data. This parameter controls a single request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.
            raw_response_callback: an optional callback invoked at each new
                response coming from the API. The only argument is the raw
                API response and the callback must return None.

        Returns:
            generator: A generator yielding documents in the paginated result set.
        """
        partialed_find = partial(
            self.find,
            filter=filter,
            projection=projection,
            sort=sort,
            timeout_info=timeout_info,
        )
        if prefetched:
            return _PrefetchIterator(
                request_method=partialed_find,
                options=options,
                prefetched=prefetched,
                raw_response_callback=raw_response_callback,
            )
        return self.paginate(
            request_method=partialed_find,
            options=options,
            raw_response_callback=raw_response_callback,
        )

    def pop(
        self,
        filter: dict[str, Any],
        pop: dict[str, Any],
        options: dict[str, Any],
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Pop the last data in the tags array

        Args:
            filter (dict): Criteria to identify the document to update.
            pop (dict): The pop to apply to the tags.
            options (dict): Additional options for the update operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The original document before the update.
        """
        json_query = make_payload(
            top_level="findOneAndUpdate",
            filter=filter,
            update={"$pop": pop},
            options=options,
        )

        response = self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def push(
        self,
        filter: dict[str, Any],
        push: dict[str, Any],
        options: dict[str, Any],
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Push new data to the tags array

        Args:
            filter (dict): Criteria to identify the document to update.
            push (dict): The push to apply to the tags.
            options (dict): Additional options for the update operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The result of the update operation.
        """
        json_query = make_payload(
            top_level="findOneAndUpdate",
            filter=filter,
            update={"$push": push},
            options=options,
        )

        response = self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def find_one_and_replace(
        self,
        replacement: dict[str, Any],
        *,
        filter: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        sort: dict[str, Any] | None = None,
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find a single document and replace it.

        Args:
            replacement (dict): The new document to replace the existing one.
            filter (dict, optional): Criteria to filter documents.
            sort (dict, optional): Specifies the order in which to find the document.
            options (dict, optional): Additional options for the operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The result of the find and replace operation.
        """
        json_query = make_payload(
            top_level="findOneAndReplace",
            filter=filter,
            projection=projection,
            replacement=replacement,
            options=options,
            sort=sort,
        )

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def vector_find_one_and_replace(
        self,
        vector: list[float],
        replacement: dict[str, Any],
        *,
        filter: dict[str, Any] | None = None,
        fields: list[str] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_DOC | None:
        """
        Perform a vector-based search and replace the first matched document.

        Args:
            vector (dict): The vector to search with.
            replacement (dict): The new document to replace the existing one.
            filter (dict, optional): Criteria to filter documents.
            fields (list, optional): Specifies the fields to return in the result.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict or None: either the matched document or None if nothing found
        """
        # Pre-process the included arguments
        sort, projection = self._recast_as_sort_projection(
            convert_vector_to_floats(vector),
            fields=fields,
        )

        # Call the underlying find() method to search
        raw_find_result = self.find_one_and_replace(
            replacement=replacement,
            filter=filter,
            projection=projection,
            sort=sort,
            timeout_info=timeout_info,
        )

        return cast(Union[API_DOC, None], raw_find_result["data"]["document"])

    def find_one_and_update(
        self,
        update: dict[str, Any],
        sort: dict[str, Any] | None = {},
        filter: dict[str, Any] | None = None,
        options: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find a single document and update it.

        Args:
            update (dict): The update to apply to the document.
            sort (dict, optional): Specifies the order in which to find the document.
            filter (dict, optional): Criteria to filter documents.
            options (dict, optional): Additional options for the operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The result of the find and update operation.
        """
        json_query = make_payload(
            top_level="findOneAndUpdate",
            filter=filter,
            update=update,
            options=options,
            sort=sort,
            projection=projection,
        )

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def vector_find_one_and_update(
        self,
        vector: list[float],
        update: dict[str, Any],
        *,
        filter: dict[str, Any] | None = None,
        fields: list[str] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_DOC | None:
        """
        Perform a vector-based search and update the first matched document.

        Args:
            vector (list): The vector to search with.
            update (dict): The update to apply to the matched document.
            filter (dict, optional): Criteria to filter documents before applying the vector search.
            fields (list, optional): Specifies the fields to return in the updated document.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict or None: The result of the vector-based find and
                update operation, or None if nothing found
        """
        # Pre-process the included arguments
        sort, projection = self._recast_as_sort_projection(
            convert_vector_to_floats(vector),
            fields=fields,
        )

        # Call the underlying find() method to search
        raw_find_result = self.find_one_and_update(
            update=update,
            filter=filter,
            sort=sort,
            projection=projection,
            timeout_info=timeout_info,
        )

        return cast(Union[API_DOC, None], raw_find_result["data"]["document"])

    def find_one_and_delete(
        self,
        sort: dict[str, Any] | None = {},
        filter: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find a single document and delete it.

        Args:
            sort (dict, optional): Specifies the order in which to find the document.
            filter (dict, optional): Criteria to filter documents.
            projection (dict, optional): Specifies the fields to return.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The result of the find and delete operation.
        """
        json_query = make_payload(
            top_level="findOneAndDelete",
            filter=filter,
            sort=sort,
            projection=projection,
        )

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def count_documents(
        self, filter: dict[str, Any] = {}, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        """
        Count documents matching a given predicate (expressed as filter).

        Args:
            filter (dict, defaults to {}): Criteria to filter documents.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: the response, either
                {"status": {"count": <NUMBER> }}
            or
                {"errors": [...]}
        """
        json_query = make_payload(
            top_level="countDocuments",
            filter=filter,
        )

        response = self._post(document=json_query, timeout_info=timeout_info)

        return response

    def find_one(
        self,
        filter: dict[str, Any] | None = {},
        projection: dict[str, Any] | None = {},
        sort: dict[str, Any] | None = {},
        options: dict[str, Any] | None = {},
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find a single document in the collection.

        Args:
            filter (dict, optional): Criteria to filter documents.
            projection (dict, optional): Specifies the fields to return.
            sort (dict, optional): Specifies the order in which to return the document.
            options (dict, optional): Additional options for the query.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: the response, either
                {"data": {"document": <DOCUMENT> }}
            or
                {"data": {"document": None}}
            depending on whether a matching document is found or not.
        """
        json_query = make_payload(
            top_level="findOne",
            filter=filter,
            projection=projection,
            options=options,
            sort=sort,
        )

        response = self._post(document=json_query, timeout_info=timeout_info)

        return response

    def vector_find_one(
        self,
        vector: list[float],
        *,
        filter: dict[str, Any] | None = None,
        fields: list[str] | None = None,
        include_similarity: bool = True,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_DOC | None:
        """
        Perform a vector-based search to find a single document in the collection.

        Args:
            vector (list): The vector to search with.
            filter (dict, optional): Additional criteria to filter documents.
            fields (list, optional): Specifies the fields to return in the result.
            include_similarity (bool, optional): Whether to include similarity score in the result.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict or None: The found document or None if no matching document is found.
        """
        # Pre-process the included arguments
        sort, projection = self._recast_as_sort_projection(
            convert_vector_to_floats(vector),
            fields=fields,
        )

        # Call the underlying find() method to search
        raw_find_result = self.find_one(
            filter=filter,
            projection=projection,
            sort=sort,
            options={"includeSimilarity": include_similarity},
            timeout_info=timeout_info,
        )

        return cast(Union[API_DOC, None], raw_find_result["data"]["document"])

    def insert_one(
        self,
        document: API_DOC,
        failures_allowed: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Insert a single document into the collection.

        Args:
            document (dict): The document to insert.
            failures_allowed (bool): Whether to allow failures in the insert operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the insert operation.
        """
        json_query = make_payload(top_level="insertOne", document=document)

        response = self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=json_query,
            skip_error_check=failures_allowed,
            timeout_info=timeout_info,
        )

        return response

    def insert_many(
        self,
        documents: list[API_DOC],
        options: dict[str, Any] | None = None,
        partial_failures_allowed: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Insert multiple documents into the collection.

        Args:
            documents (list): A list of documents to insert.
            options (dict, optional): Additional options for the insert operation.
            partial_failures_allowed (bool, optional): Whether to allow partial
                failures through the insertion (i.e. on some documents).
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the insert operation.
        """
        json_query = make_payload(
            top_level="insertMany", documents=documents, options=options
        )

        # Send the data
        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            skip_error_check=partial_failures_allowed,
            timeout_info=timeout_info,
        )

        return response

    def chunked_insert_many(
        self,
        documents: list[API_DOC],
        options: dict[str, Any] | None = None,
        partial_failures_allowed: bool = False,
        chunk_size: int = DEFAULT_INSERT_NUM_DOCUMENTS,
        concurrency: int = 1,
        timeout_info: TimeoutInfoWideType = None,
    ) -> list[API_RESPONSE | Exception]:
        """
        Insert multiple documents into the collection, handling chunking and
        optionally with concurrent insertions.

        Args:
            documents (list): A list of documents to insert.
            options (dict, optional): Additional options for the insert operation.
            partial_failures_allowed (bool, optional): Whether to allow partial
                failures in the chunk. Should be used combined with
                options={"ordered": False} in most cases.
            chunk_size (int, optional): Override the default insertion chunk size.
            concurrency (int, optional): The number of concurrent chunk insertions.
                Default is no concurrency.
            timeout_info: a float, or a TimeoutInfo dict, for each single HTTP request.
                This method runs a number of HTTP requests as it works on chunked
                data. The timeout refers to each individual such request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            list: The responses from the database after the chunked insert operation.
                This is a list of individual responses from the API: the caller
                will need to inspect them all, e.g. to collate the inserted IDs.
        """
        results: list[API_RESPONSE | Exception] = []

        # Raise a warning if ordered and concurrency
        if options and options.get("ordered") is True and concurrency > 1:
            logger.warning(
                "Using ordered insert with concurrency may lead to unexpected results."
            )

        # If we have concurrency as 1, don't use a thread pool
        if concurrency == 1:
            # Split the documents into chunks
            for i in range(0, len(documents), chunk_size):
                try:
                    results.append(
                        self.insert_many(
                            documents[i : i + chunk_size],
                            options,
                            partial_failures_allowed,
                            timeout_info=timeout_info,
                        )
                    )
                except APIRequestError as e:
                    if partial_failures_allowed:
                        results.append(e)
                    else:
                        raise e
            return results

        # Perform the bulk insert with concurrency otherwise
        with ThreadPoolExecutor(max_workers=concurrency) as executor:
            # Submit the jobs
            futures = [
                executor.submit(
                    self.insert_many,
                    documents[i : i + chunk_size],
                    options,
                    partial_failures_allowed,
                    timeout_info=timeout_info,
                )
                for i in range(0, len(documents), chunk_size)
            ]

            # Collect the results
            for future in futures:
                try:
                    results.append(future.result())
                except APIRequestError as e:
                    if partial_failures_allowed:
                        results.append(e)
                    else:
                        raise e

        return results

    def update_one(
        self,
        filter: dict[str, Any],
        update: dict[str, Any],
        sort: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
        options: dict[str, Any] | None = None,
    ) -> API_RESPONSE:
        """
        Update a single document in the collection.

        Args:
            filter (dict): Criteria to identify the document to update.
            update (dict): The update to apply to the document.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.
            options (dict, optional): Additional options for the operation.

        Returns:
            dict: The response from the database after the update operation.
        """
        json_query = make_payload(
            top_level="updateOne",
            filter=filter,
            update=update,
            options=options,
            sort=sort,
        )

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def update_many(
        self,
        filter: dict[str, Any],
        update: dict[str, Any],
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Updates multiple documents in the collection.

        Args:
            filter (dict): Criteria to identify the document to update.
            update (dict): The update to apply to the document.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the update operation.
        """
        json_query = make_payload(
            top_level="updateMany",
            filter=filter,
            update=update,
            options=options,
        )

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def replace(
        self, path: str, document: API_DOC, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        """
        Replace a document in the collection.

        Args:
            path (str): The path to the document to replace.
            document (dict): The new document to replace the existing one.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the replace operation.
        """
        return self._put(path=path, document=document, timeout_info=timeout_info)

    def delete_one(
        self,
        id: str,
        sort: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Delete a single document from the collection based on its ID.

        Args:
            id (str): The ID of the document to delete.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the delete operation.
        """
        json_query = make_payload(
            top_level="deleteOne",
            filter={"_id": id},
            sort=sort,
        )

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def delete_one_by_predicate(
        self,
        filter: dict[str, Any],
        sort: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Delete a single document from the collection based on a filter clause

        Args:
            filter: any filter dictionary
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the delete operation.
        """
        json_query = make_payload(
            top_level="deleteOne",
            filter=filter,
            sort=sort,
        )

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def delete_many(
        self,
        filter: dict[str, Any],
        skip_error_check: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Delete many documents from the collection based on a filter condition

        Args:
            filter (dict): Criteria to identify the documents to delete.
            skip_error_check (bool): whether to ignore the check for API error
                and return the response untouched. Default is False.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the delete operation.
        """
        json_query = {
            "deleteMany": {
                "filter": filter,
            }
        }

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            skip_error_check=skip_error_check,
            timeout_info=timeout_info,
        )

        return response

    def chunked_delete_many(
        self, filter: dict[str, Any], timeout_info: TimeoutInfoWideType = None
    ) -> list[API_RESPONSE]:
        """
        Delete many documents from the collection based on a filter condition,
        chaining several API calls until exhaustion of the documents to delete.

        Args:
            filter (dict): Criteria to identify the documents to delete.
            timeout_info: a float, or a TimeoutInfo dict, for each single HTTP request.
                This method runs a number of HTTP requests as it works on a
                pagination basis. The timeout refers to each individual such request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            List[dict]: The responses from the database from all the calls
        """
        responses = []
        must_proceed = True
        while must_proceed:
            dm_response = self.delete_many(filter=filter, timeout_info=timeout_info)
            responses.append(dm_response)
            must_proceed = dm_response.get("status", {}).get("moreData", False)
        return responses

    def clear(self, timeout_info: TimeoutInfoWideType = None) -> API_RESPONSE:
        """
        Clear the collection, deleting all documents

        Args:
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database.
        """
        clear_response = self.delete_many(filter={}, timeout_info=timeout_info)

        if clear_response.get("status", {}).get("deletedCount") != -1:
            raise ValueError(
                f"Could not issue a clear-collection API command (response: {json.dumps(clear_response)})."
            )

        return clear_response

    def delete_subdocument(
        self, id: str, subdoc: str, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        """
        Delete a subdocument or field from a document in the collection.

        Args:
            id (str): The ID of the document containing the subdocument.
            subdoc (str): The key of the subdocument or field to remove.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the update operation.
        """
        json_query = {
            "findOneAndUpdate": {
                "filter": {"_id": id},
                "update": {"$unset": {subdoc: ""}},
            }
        }

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def upsert_one(
        self, document: API_DOC, timeout_info: TimeoutInfoWideType = None
    ) -> str:
        """
        Emulate an upsert operation for a single document in the collection.

        This method attempts to insert the document.
        If a document with the same _id exists, it updates the existing document.

        Args:
            document (dict): The document to insert or update.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP requests.
                This method may issue one or two requests, depending on what
                is detected on DB. This timeout controls each HTTP request individually.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            str: The _id of the inserted or updated document.
        """
        # Build the payload for the insert attempt
        result = self.insert_one(
            document, failures_allowed=True, timeout_info=timeout_info
        )

        # If the call failed because of preexisting doc, then we replace it
        if "errors" in result:
            if (
                "errorCode" in result["errors"][0]
                and result["errors"][0]["errorCode"] == "DOCUMENT_ALREADY_EXISTS"
            ):
                # Now we attempt the update
                result = self.find_one_and_replace(
                    replacement=document,
                    filter={"_id": document["_id"]},
                    timeout_info=timeout_info,
                )
                upserted_id = cast(str, result["data"]["document"]["_id"])
            else:
                raise ValueError(result)
        else:
            if result.get("status", {}).get("insertedIds", []):
                upserted_id = cast(str, result["status"]["insertedIds"][0])
            else:
                raise ValueError("Unexplained empty insertedIds from API")

        return upserted_id

    def upsert_many(
        self,
        documents: list[API_DOC],
        concurrency: int = 1,
        partial_failures_allowed: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> list[str | Exception]:
        """
        Emulate an upsert operation for multiple documents in the collection.

        This method attempts to insert the documents.
        If a document with the same _id exists, it updates the existing document.

        Args:
            documents (List[dict]): The documents to insert or update.
            concurrency (int, optional): The number of concurrent upserts.
            partial_failures_allowed (bool, optional): Whether to allow partial
                failures in the batch.
            timeout_info: a float, or a TimeoutInfo dict, for each HTTP request.
                This method issues a separate HTTP request for each document to
                insert: the timeout controls each such request individually.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            List[Union[str, Exception]]: A list of "_id"s of the inserted or updated documents.
        """
        results: list[str | Exception] = []

        # If concurrency is 1, no need for thread pool
        if concurrency == 1:
            for document in documents:
                try:
                    results.append(self.upsert_one(document, timeout_info=timeout_info))
                except Exception as e:
                    results.append(e)
            return results

        # Perform the bulk upsert with concurrency
        with ThreadPoolExecutor(max_workers=concurrency) as executor:
            # Submit the jobs
            futures = [
                executor.submit(self.upsert_one, document, timeout_info=timeout_info)
                for document in documents
            ]

            # Collect the results
            for future in futures:
                try:
                    results.append(future.result())
                except Exception as e:
                    if partial_failures_allowed:
                        results.append(e)
                    else:
                        raise e

        return results


class AsyncAstraDBCollection:
    def __init__(
        self,
        collection_name: str,
        astra_db: AsyncAstraDB | None = None,
        token: str | None = None,
        api_endpoint: str | None = None,
        namespace: str | None = None,
        caller_name: str | None = None,
        caller_version: str | None = None,
        additional_headers: dict[str, str] = {},
    ) -> None:
        """
        Initialize an AstraDBCollection instance.

        Args:
            collection_name (str): The name of the collection.
            astra_db (AstraDB, optional): An instance of Astra DB.
            token (str, optional): Authentication token for Astra DB.
            api_endpoint (str, optional): API endpoint URL.
            namespace (str, optional): Namespace for the database.
            caller_name (str, optional): identity of the caller ("my_framework")
                If passing a client, its caller is used as fallback
            caller_version (str, optional): version of the caller code ("1.0.3")
                If passing a client, its caller is used as fallback
            additional_headers (Dict[str, str]): any further set of headers,
                in the form of key-value pairs, to be passed with the HTTP
                requests by this collection instance.
        """
        # Check for presence of the Astra DB object
        if astra_db is None:
            if api_endpoint is None:
                raise AssertionError("Must provide api_endpoint")

            astra_db = AsyncAstraDB(
                token=token,
                api_endpoint=api_endpoint,
                namespace=namespace,
                caller_name=caller_name,
                caller_version=caller_version,
            )
        else:
            # if astra_db passed, copy and apply possible overrides
            astra_db = astra_db.copy(
                token=token,
                api_endpoint=api_endpoint,
                namespace=namespace,
                caller_name=caller_name,
                caller_version=caller_version,
            )

        # Set the remaining instance attributes
        self.astra_db: AsyncAstraDB = astra_db
        self.caller_name: str | None = self.astra_db.caller_name
        self.caller_version: str | None = self.astra_db.caller_version
        self.additional_headers = additional_headers
        self.client = astra_db.client
        self.collection_name = collection_name
        self.base_path: str = f"{self.astra_db.base_path}/{self.collection_name}"

    def __repr__(self) -> str:
        return f'AsyncAstraDBCollection[astra_db="{self.astra_db}", collection_name="{self.collection_name}"]'

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, AsyncAstraDBCollection):
            return all(
                [
                    self.collection_name == other.collection_name,
                    self.astra_db == other.astra_db,
                    self.caller_name == other.caller_name,
                    self.caller_version == other.caller_version,
                    self.additional_headers == other.additional_headers,
                ]
            )
        else:
            return False

    def copy(
        self,
        *,
        collection_name: str | None = None,
        token: str | None = None,
        api_endpoint: str | None = None,
        api_path: str | None = None,
        api_version: str | None = None,
        namespace: str | None = None,
        caller_name: str | None = None,
        caller_version: str | None = None,
        additional_headers: dict[str, str] | None = None,
    ) -> AsyncAstraDBCollection:
        return AsyncAstraDBCollection(
            collection_name=collection_name or self.collection_name,
            astra_db=self.astra_db.copy(
                token=token,
                api_endpoint=api_endpoint,
                api_path=api_path,
                api_version=api_version,
                namespace=namespace,
                caller_name=caller_name,
                caller_version=caller_version,
            ),
            caller_name=caller_name or self.caller_name,
            caller_version=caller_version or self.caller_version,
            additional_headers=additional_headers or self.additional_headers,
        )

    def set_caller(
        self,
        caller_name: str | None = None,
        caller_version: str | None = None,
    ) -> None:
        self.astra_db.set_caller(
            caller_name=caller_name,
            caller_version=caller_version,
        )
        self.caller_name = caller_name
        self.caller_version = caller_version

    def to_sync(self) -> AstraDBCollection:
        return AstraDBCollection(
            collection_name=self.collection_name,
            astra_db=self.astra_db.to_sync(),
            caller_name=self.caller_name,
            caller_version=self.caller_version,
            additional_headers=self.additional_headers,
        )

    async def _request(
        self,
        method: str = http_methods.POST,
        path: str | None = None,
        json_data: dict[str, Any] | None = None,
        url_params: dict[str, Any] | None = None,
        skip_error_check: bool = False,
        timeout_info: TimeoutInfoWideType = None,
        **kwargs: Any,
    ) -> API_RESPONSE:
        adirect_response = await async_api_request(
            client=self.client,
            base_url=self.astra_db.base_url,
            auth_header=DEFAULT_AUTH_HEADER,
            token=self.astra_db.token,
            method=method,
            json_data=normalize_for_api(json_data),
            url_params=url_params,
            path=path,
            skip_error_check=skip_error_check,
            caller_name=self.caller_name,
            caller_version=self.caller_version,
            timeout=to_httpx_timeout(timeout_info),
            additional_headers=self.additional_headers,
        )
        response = restore_from_api(adirect_response)
        return response

    async def post_raw_request(
        self, body: dict[str, Any], timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        return await self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=body,
            timeout_info=timeout_info,
        )

    async def _get(
        self,
        path: str | None = None,
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE | None:
        full_path = f"{self.base_path}/{path}" if path else self.base_path
        response = await self._request(
            method=http_methods.GET,
            path=full_path,
            url_params=options,
            timeout_info=timeout_info,
        )
        if isinstance(response, dict):
            return response
        return None

    async def _put(
        self,
        path: str | None = None,
        document: API_RESPONSE | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        full_path = f"{self.base_path}/{path}" if path else self.base_path
        response = await self._request(
            method=http_methods.PUT,
            path=full_path,
            json_data=document,
            timeout_info=timeout_info,
        )
        return response

    async def _post(
        self,
        path: str | None = None,
        document: API_DOC | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        full_path = f"{self.base_path}/{path}" if path else self.base_path
        response = await self._request(
            method=http_methods.POST,
            path=full_path,
            json_data=document,
            timeout_info=timeout_info,
        )
        return response

    def _recast_as_sort_projection(
        self, vector: list[float], fields: list[str] | None = None
    ) -> tuple[dict[str, Any], dict[str, Any] | None]:
        """
        Given a vector and optionally a list of fields,
        reformulate them as a sort, projection pair for regular
        'find'-like API calls (with basic validation as well).
        """
        # Must pass a vector
        if not vector:
            raise ValueError("Must pass a vector")

        # Edge case for field selection
        if fields and "$similarity" in fields:
            raise ValueError("Please use the `include_similarity` parameter")

        # Build the new vector parameter
        sort: dict[str, Any] = {"$vector": vector}

        # Build the new fields parameter
        # Note: do not leave projection={}, make it None
        # (or it will devour $similarity away in the API response)
        if fields is not None and len(fields) > 0:
            projection = {f: 1 for f in fields}
        else:
            projection = None

        return sort, projection

    async def get(
        self, path: str | None = None, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE | None:
        """
        Retrieve a document from the collection by its path.

        Args:
            path (str, optional): The path of the document to retrieve.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The retrieved document.
        """
        return await self._get(path=path, timeout_info=timeout_info)

    async def find(
        self,
        filter: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        sort: dict[str, Any] | None = None,
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find documents in the collection that match the given filter.

        Args:
            filter (dict, optional): Criteria to filter documents.
            projection (dict, optional): Specifies the fields to return.
            sort (dict, optional): Specifies the order in which to return matching documents.
            options (dict, optional): Additional options for the query.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The query response containing matched documents.
        """
        json_query = make_payload(
            top_level="find",
            filter=filter,
            projection=projection,
            options=options,
            sort=sort,
        )

        response = await self._post(document=json_query, timeout_info=timeout_info)

        return response

    async def vector_find(
        self,
        vector: list[float],
        *,
        limit: int,
        filter: dict[str, Any] | None = None,
        fields: list[str] | None = None,
        include_similarity: bool = True,
        timeout_info: TimeoutInfoWideType = None,
    ) -> list[API_DOC]:
        """
        Perform a vector-based search in the collection.

        Args:
            vector (list): The vector to search with.
            limit (int): The maximum number of documents to return.
            filter (dict, optional): Criteria to filter documents.
            fields (list, optional): Specifies the fields to return.
            include_similarity (bool, optional): Whether to include similarity score in the result.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            list: A list of documents matching the vector search criteria.
        """
        # Must pass a limit
        if not limit:
            raise ValueError("Must pass a limit")

        # Pre-process the included arguments
        sort, projection = self._recast_as_sort_projection(
            vector,
            fields=fields,
        )

        # Call the underlying find() method to search
        raw_find_result = await self.find(
            filter=filter,
            projection=projection,
            sort=sort,
            options={
                "limit": limit,
                "includeSimilarity": include_similarity,
            },
            timeout_info=timeout_info,
        )

        return cast(List[API_DOC], raw_find_result["data"]["documents"])

    @staticmethod
    async def paginate(
        *,
        request_method: AsyncPaginableRequestMethod,
        options: dict[str, Any] | None,
        prefetched: int | None = None,
        timeout_info: TimeoutInfoWideType = None,
        raw_response_callback: Callable[[dict[str, Any]], None] | None = None,
    ) -> AsyncGenerator[API_DOC, None]:
        """
        Generate paginated results for a given database query method.

        Args:
            request_method (function): The database query method to paginate.
            options (dict, optional): Options for the database query.
            prefetched (int, optional): Number of pre-fetched documents.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.
            raw_response_callback: an optional callback invoked at each new
                response coming from the API. The only argument is the raw
                API response and the callback must return None.

        Yields:
            dict: The next document in the paginated result set.
        """
        _options = options or {}
        response0 = await request_method(options=_options)
        if raw_response_callback:
            raw_response_callback(response0)
        next_page_state = response0["data"]["nextPageState"]
        options0 = _options
        if next_page_state is not None and prefetched:

            async def queued_paginate(
                queue: asyncio.Queue[API_DOC | None],
                request_method: AsyncPaginableRequestMethod,
                options: dict[str, Any] | None,
            ) -> None:
                try:
                    async for doc in AsyncAstraDBCollection.paginate(
                        request_method=request_method, options=options
                    ):
                        await queue.put(doc)
                finally:
                    await queue.put(None)

            queue: asyncio.Queue[API_DOC | None] = asyncio.Queue(prefetched)
            options1 = {**options0, **{"pageState": next_page_state}}
            asyncio.create_task(queued_paginate(queue, request_method, options1))
            for document in response0["data"]["documents"]:
                yield document
            doc = await queue.get()
            while doc is not None:
                yield doc
                doc = await queue.get()
        else:
            for document in response0["data"]["documents"]:
                yield document
            while next_page_state is not None:
                options1 = {**options0, **{"pageState": next_page_state}}
                response1 = await request_method(options=options1)
                if raw_response_callback:
                    raw_response_callback(response1)
                for document in response1["data"]["documents"]:
                    yield document
                next_page_state = response1["data"]["nextPageState"]

    def paginated_find(
        self,
        filter: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        sort: dict[str, Any] | None = None,
        options: dict[str, Any] | None = None,
        prefetched: int | None = None,
        timeout_info: TimeoutInfoWideType = None,
        raw_response_callback: Callable[[dict[str, Any]], None] | None = None,
    ) -> AsyncIterator[API_DOC]:
        """
        Perform a paginated search in the collection.

        Args:
            filter (dict, optional): Criteria to filter documents.
            projection (dict, optional): Specifies the fields to return.
            sort (dict, optional): Specifies the order in which to return matching documents.
            options (dict, optional): Additional options for the query.
            prefetched (int, optional): Number of pre-fetched documents
            timeout_info: a float, or a TimeoutInfo dict, for each
                single HTTP request.
                This is a paginated method, that issues several requests as it
                needs more data. This parameter controls a single request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.
            raw_response_callback: an optional callback invoked at each new
                response coming from the API. The only argument is the raw
                API response and the callback must return None.

        Returns:
            generator: A generator yielding documents in the paginated result set.
        """
        partialed_find = partial(
            self.find,
            filter=filter,
            projection=projection,
            sort=sort,
            timeout_info=timeout_info,
        )
        return self.paginate(
            request_method=partialed_find,
            options=options,
            prefetched=prefetched,
            raw_response_callback=raw_response_callback,
        )

    async def pop(
        self,
        filter: dict[str, Any],
        pop: dict[str, Any],
        options: dict[str, Any],
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Pop the last data in the tags array

        Args:
            filter (dict): Criteria to identify the document to update.
            pop (dict): The pop to apply to the tags.
            options (dict): Additional options for the update operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The original document before the update.
        """
        json_query = make_payload(
            top_level="findOneAndUpdate",
            filter=filter,
            update={"$pop": pop},
            options=options,
        )

        response = await self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def push(
        self,
        filter: dict[str, Any],
        push: dict[str, Any],
        options: dict[str, Any],
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Push new data to the tags array

        Args:
            filter (dict): Criteria to identify the document to update.
            push (dict): The push to apply to the tags.
            options (dict): Additional options for the update operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The result of the update operation.
        """
        json_query = make_payload(
            top_level="findOneAndUpdate",
            filter=filter,
            update={"$push": push},
            options=options,
        )

        response = await self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def find_one_and_replace(
        self,
        replacement: dict[str, Any],
        *,
        filter: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        sort: dict[str, Any] | None = None,
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find a single document and replace it.

        Args:
            replacement (dict): The new document to replace the existing one.
            filter (dict, optional): Criteria to filter documents.
            sort (dict, optional): Specifies the order in which to find the document.
            options (dict, optional): Additional options for the operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The result of the find and replace operation.
        """
        json_query = make_payload(
            top_level="findOneAndReplace",
            filter=filter,
            projection=projection,
            replacement=replacement,
            options=options,
            sort=sort,
        )

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def vector_find_one_and_replace(
        self,
        vector: list[float],
        replacement: dict[str, Any],
        *,
        filter: dict[str, Any] | None = None,
        fields: list[str] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_DOC | None:
        """
        Perform a vector-based search and replace the first matched document.

        Args:
            vector (dict): The vector to search with.
            replacement (dict): The new document to replace the existing one.
            filter (dict, optional): Criteria to filter documents.
            fields (list, optional): Specifies the fields to return in the result.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict or None: either the matched document or None if nothing found
        """
        # Pre-process the included arguments
        sort, projection = self._recast_as_sort_projection(
            vector,
            fields=fields,
        )

        # Call the underlying find() method to search
        raw_find_result = await self.find_one_and_replace(
            replacement=replacement,
            filter=filter,
            projection=projection,
            sort=sort,
            timeout_info=timeout_info,
        )

        return cast(Union[API_DOC, None], raw_find_result["data"]["document"])

    async def find_one_and_update(
        self,
        update: dict[str, Any],
        sort: dict[str, Any] | None = {},
        filter: dict[str, Any] | None = None,
        options: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find a single document and update it.

        Args:
            sort (dict, optional): Specifies the order in which to find the document.
            update (dict): The update to apply to the document.
            filter (dict, optional): Criteria to filter documents.
            options (dict, optional): Additional options for the operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The result of the find and update operation.
        """
        json_query = make_payload(
            top_level="findOneAndUpdate",
            filter=filter,
            update=update,
            options=options,
            sort=sort,
            projection=projection,
        )

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def vector_find_one_and_update(
        self,
        vector: list[float],
        update: dict[str, Any],
        *,
        filter: dict[str, Any] | None = None,
        fields: list[str] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_DOC | None:
        """
        Perform a vector-based search and update the first matched document.

        Args:
            vector (list): The vector to search with.
            update (dict): The update to apply to the matched document.
            filter (dict, optional): Criteria to filter documents before applying the vector search.
            fields (list, optional): Specifies the fields to return in the updated document.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict or None: The result of the vector-based find and
                update operation, or None if nothing found
        """
        # Pre-process the included arguments
        sort, projection = self._recast_as_sort_projection(
            vector,
            fields=fields,
        )

        # Call the underlying find() method to search
        raw_find_result = await self.find_one_and_update(
            update=update,
            filter=filter,
            sort=sort,
            projection=projection,
            timeout_info=timeout_info,
        )

        return cast(Union[API_DOC, None], raw_find_result["data"]["document"])

    async def find_one_and_delete(
        self,
        sort: dict[str, Any] | None = {},
        filter: dict[str, Any] | None = None,
        projection: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find a single document and delete it.

        Args:
            sort (dict, optional): Specifies the order in which to find the document.
            filter (dict, optional): Criteria to filter documents.
            projection (dict, optional): Specifies the fields to return.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The result of the find and delete operation.
        """
        json_query = make_payload(
            top_level="findOneAndDelete",
            filter=filter,
            sort=sort,
            projection=projection,
        )

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def count_documents(
        self, filter: dict[str, Any] = {}, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        """
        Count documents matching a given predicate (expressed as filter).

        Args:
            filter (dict, defaults to {}): Criteria to filter documents.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: the response, either
                {"status": {"count": <NUMBER> }}
            or
                {"errors": [...]}
        """
        json_query = make_payload(
            top_level="countDocuments",
            filter=filter,
        )

        response = await self._post(document=json_query, timeout_info=timeout_info)

        return response

    async def find_one(
        self,
        filter: dict[str, Any] | None = {},
        projection: dict[str, Any] | None = {},
        sort: dict[str, Any] | None = {},
        options: dict[str, Any] | None = {},
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Find a single document in the collection.

        Args:
            filter (dict, optional): Criteria to filter documents.
            projection (dict, optional): Specifies the fields to return.
            sort (dict, optional): Specifies the order in which to return the document.
            options (dict, optional): Additional options for the query.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: the response, either
                {"data": {"document": <DOCUMENT> }}
            or
                {"data": {"document": None}}
            depending on whether a matching document is found or not.
        """
        json_query = make_payload(
            top_level="findOne",
            filter=filter,
            projection=projection,
            options=options,
            sort=sort,
        )

        response = await self._post(document=json_query, timeout_info=timeout_info)

        return response

    async def vector_find_one(
        self,
        vector: list[float],
        *,
        filter: dict[str, Any] | None = None,
        fields: list[str] | None = None,
        include_similarity: bool = True,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_DOC | None:
        """
        Perform a vector-based search to find a single document in the collection.

        Args:
            vector (list): The vector to search with.
            filter (dict, optional): Additional criteria to filter documents.
            fields (list, optional): Specifies the fields to return in the result.
            include_similarity (bool, optional): Whether to include similarity score in the result.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict or None: The found document or None if no matching document is found.
        """
        # Pre-process the included arguments
        sort, projection = self._recast_as_sort_projection(
            vector,
            fields=fields,
        )

        # Call the underlying find() method to search
        raw_find_result = await self.find_one(
            filter=filter,
            projection=projection,
            sort=sort,
            options={"includeSimilarity": include_similarity},
            timeout_info=timeout_info,
        )

        return cast(Union[API_DOC, None], raw_find_result["data"]["document"])

    async def insert_one(
        self,
        document: API_DOC,
        failures_allowed: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Insert a single document into the collection.

        Args:
            document (dict): The document to insert.
            failures_allowed (bool): Whether to allow failures in the insert operation.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the insert operation.
        """
        json_query = make_payload(top_level="insertOne", document=document)

        response = await self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=json_query,
            skip_error_check=failures_allowed,
            timeout_info=timeout_info,
        )

        return response

    async def insert_many(
        self,
        documents: list[API_DOC],
        options: dict[str, Any] | None = None,
        partial_failures_allowed: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Insert multiple documents into the collection.

        Args:
            documents (list): A list of documents to insert.
            options (dict, optional): Additional options for the insert operation.
            partial_failures_allowed (bool, optional): Whether to allow partial
                failures through the insertion (i.e. on some documents).
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the insert operation.
        """
        json_query = make_payload(
            top_level="insertMany", documents=documents, options=options
        )

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            skip_error_check=partial_failures_allowed,
            timeout_info=timeout_info,
        )

        return response

    async def chunked_insert_many(
        self,
        documents: list[API_DOC],
        options: dict[str, Any] | None = None,
        partial_failures_allowed: bool = False,
        chunk_size: int = DEFAULT_INSERT_NUM_DOCUMENTS,
        concurrency: int = 1,
        timeout_info: TimeoutInfoWideType = None,
    ) -> list[API_RESPONSE | Exception]:
        """
        Insert multiple documents into the collection, handling chunking and
        optionally with concurrent insertions.

        Args:
            documents (list): A list of documents to insert.
            options (dict, optional): Additional options for the insert operation.
            partial_failures_allowed (bool, optional): Whether to allow partial
                failures in the chunk. Should be used combined with
                options={"ordered": False} in most cases.
            chunk_size (int, optional): Override the default insertion chunk size.
            concurrency (int, optional): The number of concurrent chunk insertions.
                Default is no concurrency.
            timeout_info: a float, or a TimeoutInfo dict, for each single HTTP request.
                This method runs a number of HTTP requests as it works on chunked
                data. The timeout refers to each individual such request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            list: The responses from the database after the chunked insert operation.
                This is a list of individual responses from the API: the caller
                will need to inspect them all, e.g. to collate the inserted IDs.
        """
        sem = asyncio.Semaphore(concurrency)

        async def concurrent_insert_many(
            docs: list[API_DOC],
            index: int,
            partial_failures_allowed: bool,
        ) -> API_RESPONSE | Exception:
            async with sem:
                logger.debug(f"Processing chunk #{index + 1} of size {len(docs)}")
                try:
                    im_result = await self.insert_many(
                        documents=docs,
                        options=options,
                        partial_failures_allowed=partial_failures_allowed,
                        timeout_info=timeout_info,
                    )
                    logger.debug(
                        f"Finished processing chunk #{index + 1} of size {len(docs)}"
                    )
                    return im_result
                except APIRequestError as e:
                    logger.debug(
                        f"Got APIRequestError while processing chunk #{index + 1} of size {len(docs)}"
                    )
                    if partial_failures_allowed:
                        return e
                    else:
                        raise e

        if concurrency > 1:
            tasks = [
                asyncio.create_task(
                    concurrent_insert_many(
                        documents[i : i + chunk_size],
                        i // chunk_size,
                        partial_failures_allowed,
                    )
                )
                for i in range(0, len(documents), chunk_size)
            ]
            results = await asyncio.gather(*tasks, return_exceptions=False)
        else:
            # this ensures the expectation of
            # "sequential strictly obeys fail-fast if ordered and concurrency==1"
            results = [
                await concurrent_insert_many(
                    documents[i : i + chunk_size],
                    i // chunk_size,
                    partial_failures_allowed,
                )
                for i in range(0, len(documents), chunk_size)
            ]
        return results

    async def update_one(
        self,
        filter: dict[str, Any],
        update: dict[str, Any],
        sort: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
        options: dict[str, Any] | None = None,
    ) -> API_RESPONSE:
        """
        Update a single document in the collection.

        Args:
            filter (dict): Criteria to identify the document to update.
            update (dict): The update to apply to the document.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.
            options (dict, optional): Additional options for the operation.

        Returns:
            dict: The response from the database after the update operation.
        """
        json_query = make_payload(
            top_level="updateOne",
            filter=filter,
            update=update,
            options=options,
            sort=sort,
        )

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def update_many(
        self,
        filter: dict[str, Any],
        update: dict[str, Any],
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Updates multiple documents in the collection.

        Args:
            filter (dict): Criteria to identify the document to update.
            update (dict): The update to apply to the document.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the update operation.
        """
        json_query = make_payload(
            top_level="updateMany",
            filter=filter,
            update=update,
            options=options,
        )

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def replace(
        self, path: str, document: API_DOC, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        """
        Replace a document in the collection.

        Args:
            path (str): The path to the document to replace.
            document (dict): The new document to replace the existing one.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the replace operation.
        """
        return await self._put(path=path, document=document, timeout_info=timeout_info)

    async def delete_one(
        self,
        id: str,
        sort: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Delete a single document from the collection based on its ID.

        Args:
            id (str): The ID of the document to delete.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the delete operation.
        """
        json_query = make_payload(
            top_level="deleteOne",
            filter={"_id": id},
            sort=sort,
        )

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def delete_one_by_predicate(
        self,
        filter: dict[str, Any],
        sort: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Delete a single document from the collection based on a filter clause

        Args:
            filter: any filter dictionary
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the delete operation.
        """
        json_query = make_payload(
            top_level="deleteOne",
            filter=filter,
            sort=sort,
        )

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def delete_many(
        self,
        filter: dict[str, Any],
        skip_error_check: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Delete many documents from the collection based on a filter condition

        Args:
            filter (dict): Criteria to identify the documents to delete.
            skip_error_check (bool): whether to ignore the check for API error
                and return the response untouched. Default is False.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the delete operation.
        """
        json_query = {
            "deleteMany": {
                "filter": filter,
            }
        }

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            skip_error_check=skip_error_check,
            timeout_info=timeout_info,
        )

        return response

    async def chunked_delete_many(
        self, filter: dict[str, Any], timeout_info: TimeoutInfoWideType = None
    ) -> list[API_RESPONSE]:
        """
        Delete many documents from the collection based on a filter condition,
        chaining several API calls until exhaustion of the documents to delete.

        Args:
            filter (dict): Criteria to identify the documents to delete.
            timeout_info: a float, or a TimeoutInfo dict, for each single HTTP request.
                This method runs a number of HTTP requests as it works on a
                pagination basis. The timeout refers to each individual such request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            List[dict]: The responses from the database from all the calls
        """
        responses = []
        must_proceed = True
        while must_proceed:
            dm_response = await self.delete_many(
                filter=filter, timeout_info=timeout_info
            )
            responses.append(dm_response)
            must_proceed = dm_response.get("status", {}).get("moreData", False)
        return responses

    async def clear(self, timeout_info: TimeoutInfoWideType = None) -> API_RESPONSE:
        """
        Clear the collection, deleting all documents

        Args:
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database.
        """
        clear_response = await self.delete_many(filter={}, timeout_info=timeout_info)

        if clear_response.get("status", {}).get("deletedCount") != -1:
            raise ValueError(
                f"Could not issue a clear-collection API command (response: {json.dumps(clear_response)})."
            )

        return clear_response

    async def delete_subdocument(
        self, id: str, subdoc: str, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        """
        Delete a subdocument or field from a document in the collection.

        Args:
            id (str): The ID of the document containing the subdocument.
            subdoc (str): The key of the subdocument or field to remove.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database after the update operation.
        """
        json_query = {
            "findOneAndUpdate": {
                "filter": {"_id": id},
                "update": {"$unset": {subdoc: ""}},
            }
        }

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def upsert_one(
        self,
        document: API_DOC,
        timeout_info: TimeoutInfoWideType = None,
    ) -> str:
        """
        Emulate an upsert operation for a single document in the collection.

        This method attempts to insert the document.
        If a document with the same _id exists, it updates the existing document.

        Args:
            document (dict): The document to insert or update.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP requests.
                This method may issue one or two requests, depending on what
                is detected on DB. This timeout controls each HTTP request individually.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            str: The _id of the inserted or updated document.
        """
        # Build the payload for the insert attempt
        result = await self.insert_one(
            document, failures_allowed=True, timeout_info=timeout_info
        )

        # If the call failed because of preexisting doc, then we replace it
        if "errors" in result:
            if (
                "errorCode" in result["errors"][0]
                and result["errors"][0]["errorCode"] == "DOCUMENT_ALREADY_EXISTS"
            ):
                # Now we attempt the update
                result = await self.find_one_and_replace(
                    replacement=document,
                    filter={"_id": document["_id"]},
                    timeout_info=timeout_info,
                )
                upserted_id = cast(str, result["data"]["document"]["_id"])
            else:
                raise ValueError(result)
        else:
            if result.get("status", {}).get("insertedIds", []):
                upserted_id = cast(str, result["status"]["insertedIds"][0])
            else:
                raise ValueError("Unexplained empty insertedIds from API")

        return upserted_id

    async def upsert_many(
        self,
        documents: list[API_DOC],
        concurrency: int = 1,
        partial_failures_allowed: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> list[str | Exception]:
        """
        Emulate an upsert operation for multiple documents in the collection.
        This method attempts to insert the documents.
        If a document with the same _id exists, it updates the existing document.

        Args:
            documents (List[dict]): The documents to insert or update.
            concurrency (int, optional): The number of concurrent upserts.
            partial_failures_allowed (bool, optional): Whether to allow partial
                failures in the batch.
            timeout_info: a float, or a TimeoutInfo dict, for each HTTP request.
                This method issues a separate HTTP request for each document to
                insert: the timeout controls each such request individually.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            List[Union[str, Exception]]: A list of "_id"s of the inserted or updated documents.
        """
        sem = asyncio.Semaphore(concurrency)

        async def concurrent_upsert(doc: API_DOC) -> str:
            async with sem:
                return await self.upsert_one(document=doc, timeout_info=timeout_info)

        tasks = [asyncio.create_task(concurrent_upsert(doc)) for doc in documents]
        results = await asyncio.gather(
            *tasks, return_exceptions=partial_failures_allowed
        )
        for result in results:
            if isinstance(result, BaseException) and not isinstance(result, Exception):
                raise result
        return results  # type: ignore


class AstraDB:
    # Initialize the shared httpx client as a class attribute
    client = httpx.Client()

    def __init__(
        self,
        token: str | None,
        api_endpoint: str,
        api_path: str | None = None,
        api_version: str | None = None,
        namespace: str | None = None,
        caller_name: str | None = None,
        caller_version: str | None = None,
    ) -> None:
        """
        Initialize an Astra DB instance.

        Args:
            token (str): Authentication token for Astra DB.
            api_endpoint (str): API endpoint URL.
            api_path (str, optional): used to override default URI construction
            api_version (str, optional): to override default URI construction
            namespace (str, optional): Namespace for the database.
            caller_name (str, optional): identity of the caller ("my_framework")
            caller_version (str, optional): version of the caller code ("1.0.3")
        """
        self.caller_name = caller_name
        self.caller_version = caller_version

        if api_endpoint is None:
            raise AssertionError("Must provide api_endpoint")

        if namespace is None:
            logger.info(
                f"ASTRA_DB_KEYSPACE is not set. Defaulting to '{DEFAULT_KEYSPACE_NAME}'"
            )
            namespace = DEFAULT_KEYSPACE_NAME

        # Store the API token
        self.token = token

        self.api_endpoint = api_endpoint

        # Set the Base URL for the API calls
        self.base_url = self.api_endpoint.strip("/")

        # Set the API version and path from the call
        self.api_path = (DEFAULT_JSON_API_PATH if api_path is None else api_path).strip(
            "/"
        )
        self.api_version = (
            DEFAULT_JSON_API_VERSION if api_version is None else api_version
        ).strip("/")

        # Set the namespace
        self.namespace = namespace

        # Finally, construct the full base path
        base_path_components = [
            comp
            for comp in (self.api_path, self.api_version, self.namespace)
            if comp != ""
        ]
        self.base_path: str = f"/{'/'.join(base_path_components)}"

    def __repr__(self) -> str:
        return f'AstraDB[endpoint="{self.base_url}", keyspace="{self.namespace}"]'

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, AstraDB):
            # work on the "normalized" quantities (stripped, etc)
            return all(
                [
                    self.token == other.token,
                    self.base_url == other.base_url,
                    self.base_path == other.base_path,
                    self.caller_name == other.caller_name,
                    self.caller_version == other.caller_version,
                ]
            )
        else:
            return False

    def copy(
        self,
        *,
        token: str | None = None,
        api_endpoint: str | None = None,
        api_path: str | None = None,
        api_version: str | None = None,
        namespace: str | None = None,
        caller_name: str | None = None,
        caller_version: str | None = None,
    ) -> AstraDB:
        return AstraDB(
            token=token or self.token,
            api_endpoint=api_endpoint or self.base_url,
            api_path=api_path or self.api_path,
            api_version=api_version or self.api_version,
            namespace=namespace or self.namespace,
            caller_name=caller_name or self.caller_name,
            caller_version=caller_version or self.caller_version,
        )

    def to_async(self) -> AsyncAstraDB:
        return AsyncAstraDB(
            token=self.token,
            api_endpoint=self.base_url,
            api_path=self.api_path,
            api_version=self.api_version,
            namespace=self.namespace,
            caller_name=self.caller_name,
            caller_version=self.caller_version,
        )

    def set_caller(
        self,
        caller_name: str | None = None,
        caller_version: str | None = None,
    ) -> None:
        self.caller_name = caller_name
        self.caller_version = caller_version

    def _request(
        self,
        method: str = http_methods.POST,
        path: str | None = None,
        json_data: dict[str, Any] | None = None,
        url_params: dict[str, Any] | None = None,
        skip_error_check: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        direct_response = api_request(
            client=self.client,
            base_url=self.base_url,
            auth_header=DEFAULT_AUTH_HEADER,
            token=self.token,
            method=method,
            json_data=normalize_for_api(json_data),
            url_params=url_params,
            path=path,
            skip_error_check=skip_error_check,
            caller_name=self.caller_name,
            caller_version=self.caller_version,
            timeout=to_httpx_timeout(timeout_info),
            additional_headers={},
        )
        response = restore_from_api(direct_response)
        return response

    def post_raw_request(
        self, body: dict[str, Any], timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        return self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=body,
            timeout_info=timeout_info,
        )

    def collection(self, collection_name: str) -> AstraDBCollection:
        """
        Retrieve a collection from the database.

        Args:
            collection_name (str): The name of the collection to retrieve.

        Returns:
            AstraDBCollection: The collection object.
        """
        return AstraDBCollection(collection_name=collection_name, astra_db=self)

    def get_collections(
        self,
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Retrieve a list of collections from the database.

        Args:
            options (dict, optional): Options to get the collection list
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: An object containing the list of collections in the database:
                {"status": {"collections": [...]}}
        """
        # Parse the options parameter
        if options is None:
            options = {}

        json_query = make_payload(
            top_level="findCollections",
            options=options,
        )

        response = self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    def create_collection(
        self,
        collection_name: str,
        *,
        options: dict[str, Any] | None = None,
        dimension: int | None = None,
        metric: str | None = None,
        service_dict: dict[str, str] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> AstraDBCollection:
        """
        Create a new collection in the database.

        Args:
            collection_name (str): The name of the collection to create.
            options (dict, optional): Options for the collection.
            dimension (int, optional): Dimension for vector search.
            metric (str, optional): Metric choice for vector search.
            service_dict (dict, optional): a definition for the $vectorize service
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            AstraDBCollection: The created collection object.
        """
        # options from named params
        vector_options = {
            k: v
            for k, v in {
                "dimension": dimension,
                "metric": metric,
                "service": service_dict,
            }.items()
            if v is not None
        }

        # overlap/merge with stuff in options.vector
        dup_params = set((options or {}).get("vector", {}).keys()) & set(
            vector_options.keys()
        )

        # If any params are duplicated, we raise an error
        if dup_params:
            dups = ", ".join(sorted(dup_params))
            raise ValueError(
                f"Parameter(s) {dups} passed both to the method and in the options"
            )

        # Build our options dictionary if we have vector options
        if vector_options:
            options = options or {}
            options["vector"] = {
                **options.get("vector", {}),
                **vector_options,
            }

        # Build the final json payload
        jsondata = {
            k: v
            for k, v in {"name": collection_name, "options": options}.items()
            if v is not None
        }

        # Make the request to the endpoint
        self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data={"createCollection": jsondata},
            timeout_info=timeout_info,
        )

        # Get the instance object as the return of the call
        return AstraDBCollection(astra_db=self, collection_name=collection_name)

    def delete_collection(
        self, collection_name: str, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        """
        Delete a collection from the database.

        Args:
            collection_name (str): The name of the collection to delete.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database.
        """
        # Make sure we provide a collection name
        if not collection_name:
            raise ValueError("Must provide a collection name")

        response = self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data={"deleteCollection": {"name": collection_name}},
            timeout_info=timeout_info,
        )

        return response


class AsyncAstraDB:
    def __init__(
        self,
        token: str | None,
        api_endpoint: str,
        api_path: str | None = None,
        api_version: str | None = None,
        namespace: str | None = None,
        caller_name: str | None = None,
        caller_version: str | None = None,
    ) -> None:
        """
        Initialize an Astra DB instance.

        Args:
            token (str): Authentication token for Astra DB.
            api_endpoint (str): API endpoint URL.
            api_path (str, optional): used to override default URI construction
            api_version (str, optional): to override default URI construction
            namespace (str, optional): Namespace for the database.
            caller_name (str, optional): identity of the caller ("my_framework")
            caller_version (str, optional): version of the caller code ("1.0.3")
        """
        self.caller_name = caller_name
        self.caller_version = caller_version

        self.client = httpx.AsyncClient()
        if api_endpoint is None:
            raise AssertionError("Must provide api_endpoint")

        if namespace is None:
            logger.info(
                f"ASTRA_DB_KEYSPACE is not set. Defaulting to '{DEFAULT_KEYSPACE_NAME}'"
            )
            namespace = DEFAULT_KEYSPACE_NAME

        # Store the API token
        self.token = token

        self.api_endpoint = api_endpoint

        # Set the Base URL for the API calls
        self.base_url = self.api_endpoint.strip("/")

        # Set the API version and path from the call
        self.api_path = (DEFAULT_JSON_API_PATH if api_path is None else api_path).strip(
            "/"
        )
        self.api_version = (
            DEFAULT_JSON_API_VERSION if api_version is None else api_version
        ).strip("/")

        # Set the namespace
        self.namespace = namespace

        # Finally, construct the full base path
        base_path_components = [
            comp
            for comp in (self.api_path, self.api_version, self.namespace)
            if comp != ""
        ]
        self.base_path: str = f"/{'/'.join(base_path_components)}"

    def __repr__(self) -> str:
        return f'AsyncAstraDB[endpoint="{self.base_url}", keyspace="{self.namespace}"]'

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, AsyncAstraDB):
            # work on the "normalized" quantities (stripped, etc)
            return all(
                [
                    self.token == other.token,
                    self.base_url == other.base_url,
                    self.base_path == other.base_path,
                    self.caller_name == other.caller_name,
                    self.caller_version == other.caller_version,
                ]
            )
        else:
            return False

    async def __aenter__(self) -> AsyncAstraDB:
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None = None,
        exc_value: BaseException | None = None,
        traceback: TracebackType | None = None,
    ) -> None:
        await self.client.aclose()

    def copy(
        self,
        *,
        token: str | None = None,
        api_endpoint: str | None = None,
        api_path: str | None = None,
        api_version: str | None = None,
        namespace: str | None = None,
        caller_name: str | None = None,
        caller_version: str | None = None,
    ) -> AsyncAstraDB:
        return AsyncAstraDB(
            token=token or self.token,
            api_endpoint=api_endpoint or self.base_url,
            api_path=api_path or self.api_path,
            api_version=api_version or self.api_version,
            namespace=namespace or self.namespace,
            caller_name=caller_name or self.caller_name,
            caller_version=caller_version or self.caller_version,
        )

    def to_sync(self) -> AstraDB:
        return AstraDB(
            token=self.token,
            api_endpoint=self.base_url,
            api_path=self.api_path,
            api_version=self.api_version,
            namespace=self.namespace,
            caller_name=self.caller_name,
            caller_version=self.caller_version,
        )

    def set_caller(
        self,
        caller_name: str | None = None,
        caller_version: str | None = None,
    ) -> None:
        self.caller_name = caller_name
        self.caller_version = caller_version

    async def _request(
        self,
        method: str = http_methods.POST,
        path: str | None = None,
        json_data: dict[str, Any] | None = None,
        url_params: dict[str, Any] | None = None,
        skip_error_check: bool = False,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        adirect_response = await async_api_request(
            client=self.client,
            base_url=self.base_url,
            auth_header=DEFAULT_AUTH_HEADER,
            token=self.token,
            method=method,
            json_data=normalize_for_api(json_data),
            url_params=url_params,
            path=path,
            skip_error_check=skip_error_check,
            caller_name=self.caller_name,
            caller_version=self.caller_version,
            timeout=to_httpx_timeout(timeout_info),
            additional_headers={},
        )
        response = restore_from_api(adirect_response)
        return response

    async def post_raw_request(
        self, body: dict[str, Any], timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        return await self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=body,
            timeout_info=timeout_info,
        )

    async def collection(self, collection_name: str) -> AsyncAstraDBCollection:
        """
        Retrieve a collection from the database.

        Args:
            collection_name (str): The name of the collection to retrieve.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            AstraDBCollection: The collection object.
        """
        return AsyncAstraDBCollection(collection_name=collection_name, astra_db=self)

    async def get_collections(
        self,
        options: dict[str, Any] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> API_RESPONSE:
        """
        Retrieve a list of collections from the database.

        Args:
            options (dict, optional): Options to get the collection list
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: An object containing the list of collections in the database:
                {"status": {"collections": [...]}}
        """
        # Parse the options parameter
        if options is None:
            options = {}

        json_query = make_payload(
            top_level="findCollections",
            options=options,
        )

        response = await self._request(
            method=http_methods.POST,
            path=self.base_path,
            json_data=json_query,
            timeout_info=timeout_info,
        )

        return response

    async def create_collection(
        self,
        collection_name: str,
        *,
        options: dict[str, Any] | None = None,
        dimension: int | None = None,
        metric: str | None = None,
        service_dict: dict[str, str] | None = None,
        timeout_info: TimeoutInfoWideType = None,
    ) -> AsyncAstraDBCollection:
        """
        Create a new collection in the database.

        Args:
            collection_name (str): The name of the collection to create.
            options (dict, optional): Options for the collection.
            dimension (int, optional): Dimension for vector search.
            metric (str, optional): Metric choice for vector search.
            service_dict (dict, optional): a definition for the $vectorize service
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            AsyncAstraDBCollection: The created collection object.
        """
        # options from named params
        vector_options = {
            k: v
            for k, v in {
                "dimension": dimension,
                "metric": metric,
                "service": service_dict,
            }.items()
            if v is not None
        }

        # overlap/merge with stuff in options.vector
        dup_params = set((options or {}).get("vector", {}).keys()) & set(
            vector_options.keys()
        )

        # If any params are duplicated, we raise an error
        if dup_params:
            dups = ", ".join(sorted(dup_params))
            raise ValueError(
                f"Parameter(s) {dups} passed both to the method and in the options"
            )

        # Build our options dictionary if we have vector options
        if vector_options:
            options = options or {}
            options["vector"] = {
                **options.get("vector", {}),
                **vector_options,
            }

        # Build the final json payload
        jsondata = {
            k: v
            for k, v in {"name": collection_name, "options": options}.items()
            if v is not None
        }

        # Make the request to the endpoint
        await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data={"createCollection": jsondata},
            timeout_info=timeout_info,
        )

        # Get the instance object as the return of the call
        return AsyncAstraDBCollection(astra_db=self, collection_name=collection_name)

    async def delete_collection(
        self, collection_name: str, timeout_info: TimeoutInfoWideType = None
    ) -> API_RESPONSE:
        """
        Delete a collection from the database.

        Args:
            collection_name (str): The name of the collection to delete.
            timeout_info: a float, or a TimeoutInfo dict, for the HTTP request.
                Note that a 'read' timeout event will not block the action taken
                by the API server if it has received the request already.

        Returns:
            dict: The response from the database.
        """
        # Make sure we provide a collection name
        if not collection_name:
            raise ValueError("Must provide a collection name")

        response = await self._request(
            method=http_methods.POST,
            path=f"{self.base_path}",
            json_data={"deleteCollection": {"name": collection_name}},
            timeout_info=timeout_info,
        )

        return response
