# Source code for limits.aio.storage.mongodb

    from __future__ import annotations

    import asyncio
    import datetime
    import time

    from deprecated.sphinx import versionadded, versionchanged

    from limits.aio.storage.base import (
        MovingWindowSupport,
        SlidingWindowCounterSupport,
        Storage,
    )
    from limits.typing import (
        ParamSpec,
        TypeVar,
        cast,
    )
    from limits.util import get_dependency

    P = ParamSpec("P")
    R = TypeVar("R")

    @versionadded(version="2.1")
    @versionchanged(
        version="3.14.0",
        reason="Added option to select custom collection names for windows & counters",
    )
    class MongoDBStorage(Storage, MovingWindowSupport, SlidingWindowCounterSupport):
        """
        Rate limit storage with MongoDB as backend.

        Depends on :pypi:`motor`
        """

        STORAGE_SCHEME = ["async+mongodb", "async+mongodb+srv"]
        """
        The storage scheme for MongoDB for use in an async context
        """

        DEPENDENCIES = ["motor.motor_asyncio", "pymongo"]

        def __init__(
            self,
            uri: str,
            database_name: str = "limits",
            counter_collection_name: str = "counters",
            window_collection_name: str = "windows",
            wrap_exceptions: bool = False,
            **options: float | str | bool,
        ) -> None:
            """
            :param uri: uri of the form ``async+mongodb://[user:password]@host:port?...``,
             This uri is passed directly to :class:`~motor.motor_asyncio.AsyncIOMotorClient`
            :param database_name: The database to use for storing the rate limit
             collections.
            :param counter_collection_name: The collection name to use for individual counters
             used in fixed window strategies
            :param window_collection_name: The collection name to use for sliding & moving window
             storage
            :param wrap_exceptions: Whether to wrap storage exceptions in
             :exc:`limits.errors.StorageError` before raising it.
            :param options: all remaining keyword arguments are passed
             to the constructor of :class:`~motor.motor_asyncio.AsyncIOMotorClient`
            :raise ConfigurationError: when the :pypi:`motor` or :pypi:`pymongo` are
             not available
            """

            uri = uri.replace("async+mongodb", "mongodb", 1)

            super().__init__(uri, wrap_exceptions=wrap_exceptions, **options)

            self.dependency = self.dependencies["motor.motor_asyncio"]
            self.proxy_dependency = self.dependencies["pymongo"]
            self.lib_errors, _ = get_dependency("pymongo.errors")

            self.storage = self.dependency.module.AsyncIOMotorClient(uri, **options)
            # TODO: Fix this hack. It was noticed when running a benchmark
            # with FastAPI - however - doesn't appear in unit tests or in an isolated
            # use. Reference: https://jira.mongodb.org/browse/MOTOR-822
            self.storage.get_io_loop = asyncio.get_running_loop

            self.__database_name = database_name
            self.__collection_mapping = {
                "counters": counter_collection_name,
                "windows": window_collection_name,
            }
            self.__indices_created = False

        @property
        def base_exceptions(
            self,
        ) -> type[Exception] | tuple[type[Exception], ...]:  # pragma: no cover
            return self.lib_errors.PyMongoError  # type: ignore

        @property
        def database(self):  # type: ignore
            return self.storage.get_database(self.__database_name)

        async def create_indices(self) -> None:
            if not self.__indices_created:
                await asyncio.gather(
                    self.database[self.__collection_mapping["counters"]].create_index(
                        "expireAt", expireAfterSeconds=0
                    ),
                    self.database[self.__collection_mapping["windows"]].create_index(
                        "expireAt", expireAfterSeconds=0
                    ),
                )
            self.__indices_created = True

        async def reset(self) -> int | None:
            """
            Delete all rate limit keys in the rate limit collections (counters, windows)
            """
            num_keys = sum(
                await asyncio.gather(
                    self.database[self.__collection_mapping["counters"]].count_documents(
                        {}
                    ),
                    self.database[self.__collection_mapping["windows"]].count_documents({}),
                )
            )
            await asyncio.gather(
                self.database[self.__collection_mapping["counters"]].drop(),
                self.database[self.__collection_mapping["windows"]].drop(),
            )

            return cast(int, num_keys)

        async def clear(self, key: str) -> None:
            """
            :param key: the key to clear rate limits for
            """
            await asyncio.gather(
                self.database[self.__collection_mapping["counters"]].find_one_and_delete(
                    {"_id": key}
                ),
                self.database[self.__collection_mapping["windows"]].find_one_and_delete(
                    {"_id": key}
                ),
            )

        async def get_expiry(self, key: str) -> float:
            """
            :param key: the key to get the expiry for
            """
            counter = await self.database[self.__collection_mapping["counters"]].find_one(
                {"_id": key}
            )
            return (
                (counter["expireAt"] if counter else datetime.datetime.now())
                .replace(tzinfo=datetime.timezone.utc)
                .timestamp()
            )

        async def get(self, key: str) -> int:
            """
            :param key: the key to get the counter value for
            """
            counter = await self.database[self.__collection_mapping["counters"]].find_one(
                {
                    "_id": key,
                    "expireAt": {"$gte": datetime.datetime.now(datetime.timezone.utc)},
                },
                projection=["count"],
            )

            return counter and counter["count"] or 0

        async def incr(
            self, key: str, expiry: int, elastic_expiry: bool = False, amount: int = 1
        ) -> int:
            """
            increments the counter for a given rate limit key

            :param key: the key to increment
            :param expiry: amount in seconds for the key to expire in
            :param elastic_expiry: whether to keep extending the rate limit
             window every hit.
            :param amount: the number to increment by
            """
            await self.create_indices()

            expiration = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
                seconds=expiry
            )

            response = await self.database[
                self.__collection_mapping["counters"]
            ].find_one_and_update(
                {"_id": key},
                [
                    {
                        "$set": {
                            "count": {
                                "$cond": {
                                    "if": {"$lt": ["$expireAt", "$$NOW"]},
                                    "then": amount,
                                    "else": {"$add": ["$count", amount]},
                                }
                            },
                            "expireAt": {
                                "$cond": {
                                    "if": {"$lt": ["$expireAt", "$$NOW"]},
                                    "then": expiration,
                                    "else": (expiration if elastic_expiry else "$expireAt"),
                                }
                            },
                        }
                    },
                ],
                upsert=True,
                projection=["count"],
                return_document=self.proxy_dependency.module.ReturnDocument.AFTER,
            )

            return int(response["count"])

        async def check(self) -> bool:
            """
            Check if storage is healthy by calling
            :meth:`motor.motor_asyncio.AsyncIOMotorClient.server_info`
            """
            try:
                await self.storage.server_info()

                return True
            except:  # noqa: E722
                return False

        async def get_moving_window(
            self, key: str, limit: int, expiry: int
        ) -> tuple[float, int]:
            """
            returns the starting point and the number of entries in the moving
            window

            :param str key: rate limit key
            :param int expiry: expiry of entry
            :return: (start of window, number of acquired entries)
            """
            timestamp = time.time()
            if result := (
                await self.database[self.__collection_mapping["windows"]]
                .aggregate(
                    [
                        {"$match": {"_id": key}},
                        {
                            "$project": {
                                "entries": {
                                    "$filter": {
                                        "input": "$entries",
                                        "as": "entry",
                                        "cond": {"$gte": ["$$entry", timestamp - expiry]},
                                    }
                                }
                            }
                        },
                        {"$unwind": "$entries"},
                        {
                            "$group": {
                                "_id": "$_id",
                                "min": {"$min": "$entries"},
                                "count": {"$sum": 1},
                            }
                        },
                    ]
                )
                .to_list(length=1)
            ):
                return result[0]["min"], result[0]["count"]
            return timestamp, 0

        async def acquire_entry(
            self, key: str, limit: int, expiry: int, amount: int = 1
        ) -> bool:
            """
            :param key: rate limit key to acquire an entry in
            :param limit: amount of entries allowed
            :param expiry: expiry of the entry
            :param amount: the number of entries to acquire
            """
            await self.create_indices()

            if amount > limit:
                return False

            timestamp = time.time()
            try:
                updates: dict[
                    str,
                    dict[str, datetime.datetime | dict[str, list[float] | int]],
                ] = {
                    "$push": {
                        "entries": {
                            "$each": [timestamp] * amount,
                            "$position": 0,
                            "$slice": limit,
                        }
                    },
                    "$set": {
                        "expireAt": (
                            datetime.datetime.now(datetime.timezone.utc)
                            + datetime.timedelta(seconds=expiry)
                        )
                    },
                }

                await self.database[self.__collection_mapping["windows"]].update_one(
                    {
                        "_id": key,
                        f"entries.{limit - amount}": {"$not": {"$gte": timestamp - expiry}},
                    },
                    updates,
                    upsert=True,
                )

                return True
            except self.proxy_dependency.module.errors.DuplicateKeyError:
                return False

        async def acquire_sliding_window_entry(
            self, key: str, limit: int, expiry: int, amount: int = 1
        ) -> bool:
            await self.create_indices()
            expiry_ms = expiry * 1000
            result = await self.database[
                self.__collection_mapping["windows"]
            ].find_one_and_update(
                {"_id": key},
                [
                    {
                        "$set": {
                            "previousCount": {
                                "$cond": {
                                    "if": {
                                        "$lte": [
                                            {"$subtract": ["$expiresAt", "$$NOW"]},
                                            expiry_ms,
                                        ]
                                    },
                                    "then": {"$ifNull": ["$currentCount", 0]},
                                    "else": {"$ifNull": ["$previousCount", 0]},
                                }
                            },
                        }
                    },
                    {
                        "$set": {
                            "currentCount": {
                                "$cond": {
                                    "if": {
                                        "$lte": [
                                            {"$subtract": ["$expiresAt", "$$NOW"]},
                                            expiry_ms,
                                        ]
                                    },
                                    "then": 0,
                                    "else": {"$ifNull": ["$currentCount", 0]},
                                }
                            },
                            "expiresAt": {
                                "$cond": {
                                    "if": {
                                        "$lte": [
                                            {"$subtract": ["$expiresAt", "$$NOW"]},
                                            expiry_ms,
                                        ]
                                    },
                                    "then": {
                                        "$cond": {
                                            "if": {"$gt": ["$expiresAt", 0]},
                                            "then": {"$add": ["$expiresAt", expiry_ms]},
                                            "else": {"$add": ["$$NOW", 2 * expiry_ms]},
                                        }
                                    },
                                    "else": "$expiresAt",
                                }
                            },
                        }
                    },
                    {
                        "$set": {
                            "curWeightedCount": {
                                "$floor": {
                                    "$add": [
                                        {
                                            "$multiply": [
                                                "$previousCount",
                                                {
                                                    "$divide": [
                                                        {
                                                            "$max": [
                                                                0,
                                                                {
                                                                    "$subtract": [
                                                                        "$expiresAt",
                                                                        {
                                                                            "$add": [
                                                                                "$$NOW",
                                                                                expiry_ms,
                                                                            ]
                                                                        },
                                                                    ]
                                                                },
                                                            ]
                                                        },
                                                        expiry_ms,
                                                    ]
                                                },
                                            ]
                                        },
                                        "$currentCount",
                                    ]
                                }
                            }
                        }
                    },
                    {
                        "$set": {
                            "currentCount": {
                                "$cond": {
                                    "if": {
                                        "$lte": [
                                            {"$add": ["$curWeightedCount", amount]},
                                            limit,
                                        ]
                                    },
                                    "then": {"$add": ["$currentCount", amount]},
                                    "else": "$currentCount",
                                }
                            }
                        }
                    },
                    {
                        "$set": {
                            "_acquired": {
                                "$lte": [{"$add": ["$curWeightedCount", amount]}, limit]
                            }
                        }
                    },
                    {"$unset": ["curWeightedCount"]},
                ],
                return_document=self.proxy_dependency.module.ReturnDocument.AFTER,
                upsert=True,
            )

            return cast(bool, result["_acquired"])

        async def get_sliding_window(
            self, key: str, expiry: int
        ) -> tuple[int, float, int, float]:
            expiry_ms = expiry * 1000
            if result := await self.database[
                self.__collection_mapping["windows"]
            ].find_one_and_update(
                {"_id": key},
                [
                    {
                        "$set": {
                            "previousCount": {
                                "$cond": {
                                    "if": {
                                        "$lte": [
                                            {"$subtract": ["$expiresAt", "$$NOW"]},
                                            expiry_ms,
                                        ]
                                    },
                                    "then": {"$ifNull": ["$currentCount", 0]},
                                    "else": {"$ifNull": ["$previousCount", 0]},
                                }
                            },
                            "currentCount": {
                                "$cond": {
                                    "if": {
                                        "$lte": [
                                            {"$subtract": ["$expiresAt", "$$NOW"]},
                                            expiry_ms,
                                        ]
                                    },
                                    "then": 0,
                                    "else": {"$ifNull": ["$currentCount", 0]},
                                }
                            },
                            "expiresAt": {
                                "$cond": {
                                    "if": {
                                        "$lte": [
                                            {"$subtract": ["$expiresAt", "$$NOW"]},
                                            expiry_ms,
                                        ]
                                    },
                                    "then": {"$add": ["$expiresAt", expiry_ms]},
                                    "else": "$expiresAt",
                                }
                            },
                        }
                    }
                ],
                return_document=self.proxy_dependency.module.ReturnDocument.AFTER,
                projection=["currentCount", "previousCount", "expiresAt"],
            ):
                expires_at = (
                    (result["expiresAt"].replace(tzinfo=datetime.timezone.utc).timestamp())
                    if result.get("expiresAt")
                    else time.time()
                )
                current_ttl = max(0, expires_at - time.time())
                prev_ttl = max(0, current_ttl - expiry if result["previousCount"] else 0)

                return (
                    result["previousCount"],
                    prev_ttl,
                    result["currentCount"],
                    current_ttl,
                )
            return 0, 0.0, 0, 0.0