"""
watch.py

Manages the watch cursor.
"""


import time
import typing

import pymongo.change_stream
import pymongo.collection
import pymongo.database
import pymongo.mongo_client

OperationType = typing.Literal["insert", "update", "delete", "replace", "drop", "rename", "dropDatabase", "invalidate"]


class Operation():
    """
    An enum for the different event that can occur on MongoDB

    Example
    -------
    >>> from yuno.direction import Operation
    >>> Operation.INSERT
    'insert'
    """
    UPDATE: OperationType = "update"
    INSERT: OperationType = "insert"
    DELETE: OperationType = "delete"
    REPLACE: OperationType = "replace"
    DROP: OperationType = "drop"
    RENAME: OperationType = "rename"
    DROP_DATABASE: OperationType = "dropDatabase"
    INVALIDATE: OperationType = "invalidate"


class WatchEvent():
    """
    An object representing an event on MongoDB.
    """
    class Namespace:
        """
        The namespace the event occured in
        """

        def __init__(self, data: dict) -> None:
            self.database = data.get("db")
            self.collection = data.get("coll")

        def __getitem__(self, key: str) -> str:
            return self.__getattribute__(key)

    class LSID:
        """
        The id of the event
        """

        def __init__(self, data: dict) -> None:
            self.id = data.get("id", None)
            self.uid = data.get("uid", None)

    def __init__(self, data: dict) -> None:
        """Initialize the object with raw event data"""
        self.id = self._id = data.get("_id")
        self.operation: OperationType = data.get("operationType")
        self.document = data.get("fullDocument", None)
        self.namespace = self.Namespace(data.get("ns", {}))
        self.timestamp = data.get("clusterTime", None)
        self.transaction = data.get("txnNumber", None)
        self.session_id = self.LSID(data.get("lsid", {}))


class CRUDEvent(WatchEvent):
    class DocumentKey:
        def __init__(self, data: dict) -> None:
            self.id = self._id = data.get("_id")
            self.__data__ = dict(data)

        def __getattribute__(self, name: str):
            if name in ("__data__", "__getitem__", "__getattribute__"):
                return super().__getattribute__(name)
            return self.__getitem__(name)

        def __getitem__(self, name: str):
            return self.data[name]

    def __init__(self, data: dict) -> None:
        super().__init__(data)
        self.document_key = self.DocumentKey(data.get("documentKey", {}))  # insert, replace, delete, update


class UpdateEvent(CRUDEvent):
    class Description:
        class Truncated:
            def __init__(self, data: dict) -> None:
                self.field = data.get("field", None)
                self.new_size = data.get("newSize", None)

        def __init__(self, data: dict) -> None:
            self.updated_fields: typing.Dict[str, typing.Any] = data.get("updatedFields", {})
            self.removed_fields: typing.List[str] = data.get("removedFields", [])
            self.truncated_arrays = [self.Truncated(e) for e in data.get("truncatedArrays", [])]

    def __init__(self, data: dict) -> None:
        super().__init__(data)
        self.update_description = self.Description(data.get("updateDescription", {}))


class InsertEvent(CRUDEvent):
    pass


class DeleteEvent(CRUDEvent):
    pass


class ReplaceEvent(CRUDEvent):
    pass


class DropEvent(WatchEvent):
    pass


class RenameEvent(WatchEvent):
    def __init__(self, data: dict) -> None:
        super().__init__(data)
        self.to = self.Namespace(data.get("to", {}))  # rename


class DropDatabaseEvent(WatchEvent):
    pass


class InvalidateEvent(WatchEvent):
    pass


class Watch():
    """
    A db.watch(), db.collection.watch(), db.collection.object.watch() stream to get the different events.
    """
    __stream__: pymongo.change_stream.ChangeStream
    __watching_object__: pymongo.collection.Collection
    __state__: dict = {
        "token": None,  # the resume token
        "time": 0,  # time of the last error
        "count": 0  # number of errors in the period
    }

    # pipeline=pipeline, full_document=None, resume_after=resume_state["token"], max_await_time_ms=None,
    #    batch_size=None, collation=None, start_at_operation_time=None, session=None, start_after=None

    def __init__(self, watching_object: typing.Union[pymongo.collection.Collection, pymongo.database.Database, pymongo.mongo_client.MongoClient], pipeline: typing.List[dict] = None, full_document: typing.Union[str, bool] = False, error_limit: int = 3, error_expiration: float = 60, **kwargs) -> None:
        """
        Initializes the stream.

        Parameters
        ----------
        watching_object: pymongo.collection.Collection, pymongo.database.Database, pymongo.mongo_client.MongoClient
            The object to watch.
        pipeline: list[dict]
            The pipeline to use.
        full_document: bool, str
            To return the full document when the event is an UpdateEvent.
        error_limit: int
            The number of errors before the stream is closed.
        error_expiration: float
            The number of seconds before the error count is reset.
        kwargs:
            The arguments to pass to the stream.
        """
        if isinstance(full_document, bool):
            full_document = "updateLookup" if full_document else None
        self.pipeline = pipeline
        self.full_document = full_document
        self.kwargs = kwargs

        self.error_limit = int(error_limit)
        self.error_expiration = float(error_expiration)

        self.__watching_object__ = watching_object
        self.__stream__ = watching_object.watch(pipeline, full_document, **kwargs)
        self.__closed__ = False

    def __next__(self):  # alias
        """
        Get the next event (blocking operation)

        Returns
        -------
        WatchEvent
            The next event to occur on the object.

        Example
        -------
        >>> document = db.collection.document
        >>> watch_obj = document.watch()
        >>> for event in watch_obj:
        >>>     print(event) # only called when an event occurs on the database
        """
        return self.next()

    def _get_right_event(self, data: dict) -> WatchEvent:
        """
        An internal function to get the right event type from raw event data.

        Parameters
        ----------
        data: dict
            The raw event data.

        Returns
        -------
        WatchEvent
            The event.
        """
        operation = data.get("operationType", None)
        if operation == Operation.UPDATE:
            return UpdateEvent(data)
        elif operation == Operation.INSERT:
            return InsertEvent(data)
        elif operation == Operation.DELETE:
            return DeleteEvent(data)
        elif operation == Operation.REPLACE:
            return ReplaceEvent(data)
        elif operation == Operation.DROP:
            return DropEvent(data)
        elif operation == Operation.RENAME:
            return RenameEvent(data)
        elif operation == Operation.DROP_DATABASE:
            return DropDatabaseEvent(data)
        elif operation == Operation.INVALIDATE:
            return InvalidateEvent(data)
        return WatchEvent(data)

    def next(self) -> WatchEvent:
        """
        Get the next event (blocking operation)

        Returns
        -------
        WatchEvent
            The next event to occur on the object.

        Example
        -------
        >>> document = db.collection.document
        >>> watch_obj = document.watch()
        >>> for event in watch_obj:
        >>>     print(event) # only called when an event occurs on the database
        """
        if self.__closed__:
            raise StopIteration("Stream is closed")
        try:
            data = self.__stream__.next()
            self.__state__["token"] = self.resume_token
        except Exception as err:
            self.__state__["count"] += 1
            if time.time() - self.__state__["time"] > self.error_expiration:
                self.__state__["time"] = time.time()
                self.__state__["count"] = 1
            else:
                self.__state__["time"] = time.time()
                if self.__state__["count"] >= self.error_limit:
                    try:
                        self.__stream__.close()
                    except Exception:
                        pass
                    raise ValueError("More than {} errors have occured in {} seconds while watching for changes in {}".format(
                        self.error_limit, self.error_expiration, self.__watching_object__)) from err
            self.kwargs["resume_after"] = self.__state__["token"]
            self.__stream__ = self.__watching_object__.watch(self.pipeline, self.full_document, **self.kwargs)
            return self.__next__()

        return self._get_right_event(data)

    def try_next(self) -> typing.Any:
        """
        Try to get the next event without raising an exception and without waiting.

        Returns
        -------
        WatchEvent
        """
        try:
            data = self.__stream__.try_next()
            if data is not None:
                return self._get_right_event(data)
            return data
        except StopIteration:
            return None

    def close(self):
        """Closes the stream."""
        self.__closed__ = True
        self.__stream__.close()

    def resume(self):
        """Resume the stream from the last known state."""
        self.kwargs["resume_after"] = self.__state__["token"]
        self.__stream__ = self.__watching_object__.watch(self.pipeline, self.full_document, **self.kwargs)
        self.__closed__ = False

    @property
    def resume_token(self):
        """Get the resume token, which is used to resume the stream if failed."""
        return self.__stream__.resume_token

    @property
    def alive(self):
        """
        Get whether the stream is alive.

        Returns
        -------
        bool
            Whether the stream is alive.
        """
        return self._cursor.alive

    def __enter__(self):
        """
        Enter the context manager.

        Example
        -------
        >>> with db.watch() as stream: # <-- this line calls __enter__
        ...     for event in stream:
        ...         print(event)
        """
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Exit the context manager.

        Example
        -------
        >>> with db.watch() as stream:
        ...     for event in stream:
        ...         print(event)
        ... # <-- this line calls __exit__
        """
        self.close()

    def __iter__(self) -> typing.Iterator[WatchEvent]:
        """
        Returns the iterator.
        """
        return self

"""
watch.py

Manages the watch cursor.
"""


import time
import typing

import pymongo.change_stream
import pymongo.collection
import pymongo.database
import pymongo.mongo_client

OperationType = typing.Literal["insert", "update", "delete", "replace", "drop", "rename", "dropDatabase", "invalidate"]


class Operation():
    """
    An enum for the different event that can occur on MongoDB

    Example
    -------
    >>> from yuno.direction import Operation
    >>> Operation.INSERT
    'insert'
    """
    UPDATE: OperationType = "update"
    INSERT: OperationType = "insert"
    DELETE: OperationType = "delete"
    REPLACE: OperationType = "replace"
    DROP: OperationType = "drop"
    RENAME: OperationType = "rename"
    DROP_DATABASE: OperationType = "dropDatabase"
    INVALIDATE: OperationType = "invalidate"


class WatchEvent():
    """
    An object representing an event on MongoDB.
    """
    class Namespace:
        """
        The namespace the event occured in
        """

        def __init__(self, data: dict) -> None:
            self.database = data.get("db")
            self.collection = data.get("coll")

        def __getitem__(self, key: str) -> str:
            return self.__getattribute__(key)

    class LSID:
        """
        The id of the event
        """

        def __init__(self, data: dict) -> None:
            self.id = data.get("id", None)
            self.uid = data.get("uid", None)

    def __init__(self, data: dict) -> None:
        """Initialize the object with raw event data"""
        self.id = self._id = data.get("_id")
        self.operation: OperationType = data.get("operationType")
        self.document = data.get("fullDocument", None)
        self.namespace = self.Namespace(data.get("ns", {}))
        self.timestamp = data.get("clusterTime", None)
        self.transaction = data.get("txnNumber", None)
        self.session_id = self.LSID(data.get("lsid", {}))


class CRUDEvent(WatchEvent):
    class DocumentKey:
        def __init__(self, data: dict) -> None:
            self.id = self._id = data.get("_id")
            self.__data__ = dict(data)

        def __getattribute__(self, name: str):
            if name in ("__data__", "__getitem__", "__getattribute__"):
                return super().__getattribute__(name)
            return self.__getitem__(name)

        def __getitem__(self, name: str):
            return self.data[name]

    def __init__(self, data: dict) -> None:
        super().__init__(data)
        self.document_key = self.DocumentKey(data.get("documentKey", {}))  # insert, replace, delete, update


class UpdateEvent(CRUDEvent):
    class Description:
        class Truncated:
            def __init__(self, data: dict) -> None:
                self.field = data.get("field", None)
                self.new_size = data.get("newSize", None)

        def __init__(self, data: dict) -> None:
            self.updated_fields: typing.Dict[str, typing.Any] = data.get("updatedFields", {})
            self.removed_fields: typing.List[str] = data.get("removedFields", [])
            self.truncated_arrays = [self.Truncated(e) for e in data.get("truncatedArrays", [])]

    def __init__(self, data: dict) -> None:
        super().__init__(data)
        self.update_description = self.Description(data.get("updateDescription", {}))


class InsertEvent(CRUDEvent):
    pass


class DeleteEvent(CRUDEvent):
    pass


class ReplaceEvent(CRUDEvent):
    pass


class DropEvent(WatchEvent):
    pass


class RenameEvent(WatchEvent):
    def __init__(self, data: dict) -> None:
        super().__init__(data)
        self.to = self.Namespace(data.get("to", {}))  # rename


class DropDatabaseEvent(WatchEvent):
    pass


class InvalidateEvent(WatchEvent):
    pass


class Watch():
    """
    A db.watch(), db.collection.watch(), db.collection.object.watch() stream to get the different events.
    """
    __stream__: pymongo.change_stream.ChangeStream
    __watching_object__: pymongo.collection.Collection
    __state__: dict = {
        "token": None,  # the resume token
        "time": 0,  # time of the last error
        "count": 0  # number of errors in the period
    }

    # pipeline=pipeline, full_document=None, resume_after=resume_state["token"], max_await_time_ms=None,
    #    batch_size=None, collation=None, start_at_operation_time=None, session=None, start_after=None

    def __init__(self, watching_object: typing.Union[pymongo.collection.Collection, pymongo.database.Database, pymongo.mongo_client.MongoClient], pipeline: typing.List[dict] = None, full_document: typing.Union[str, bool] = False, error_limit: int = 3, error_expiration: float = 60, **kwargs) -> None:
        """
        Initializes the stream.

        Parameters
        ----------
        watching_object: pymongo.collection.Collection, pymongo.database.Database, pymongo.mongo_client.MongoClient
            The object to watch.
        pipeline: list[dict]
            The pipeline to use.
        full_document: bool, str
            To return the full document when the event is an UpdateEvent.
        error_limit: int
            The number of errors before the stream is closed.
        error_expiration: float
            The number of seconds before the error count is reset.
        kwargs:
            The arguments to pass to the stream.
        """
        if isinstance(full_document, bool):
            full_document = "updateLookup" if full_document else None
        self.pipeline = pipeline
        self.full_document = full_document
        self.kwargs = kwargs

        self.error_limit = int(error_limit)
        self.error_expiration = float(error_expiration)

        self.__watching_object__ = watching_object
        self.__stream__ = watching_object.watch(pipeline, full_document, **kwargs)
        self.__closed__ = False

    def __next__(self):  # alias
        """
        Get the next event (blocking operation)

        Returns
        -------
        WatchEvent
            The next event to occur on the object.

        Example
        -------
        >>> document = db.collection.document
        >>> watch_obj = document.watch()
        >>> for event in watch_obj:
        >>>     print(event) # only called when an event occurs on the database
        """
        return self.next()

    def _get_right_event(self, data: dict) -> WatchEvent:
        """
        An internal function to get the right event type from raw event data.

        Parameters
        ----------
        data: dict
            The raw event data.

        Returns
        -------
        WatchEvent
            The event.
        """
        operation = data.get("operationType", None)
        if operation == Operation.UPDATE:
            return UpdateEvent(data)
        elif operation == Operation.INSERT:
            return InsertEvent(data)
        elif operation == Operation.DELETE:
            return DeleteEvent(data)
        elif operation == Operation.REPLACE:
            return ReplaceEvent(data)
        elif operation == Operation.DROP:
            return DropEvent(data)
        elif operation == Operation.RENAME:
            return RenameEvent(data)
        elif operation == Operation.DROP_DATABASE:
            return DropDatabaseEvent(data)
        elif operation == Operation.INVALIDATE:
            return InvalidateEvent(data)
        return WatchEvent(data)

    def next(self) -> WatchEvent:
        """
        Get the next event (blocking operation)

        Returns
        -------
        WatchEvent
            The next event to occur on the object.

        Example
        -------
        >>> document = db.collection.document
        >>> watch_obj = document.watch()
        >>> for event in watch_obj:
        >>>     print(event) # only called when an event occurs on the database
        """
        if self.__closed__:
            raise StopIteration("Stream is closed")
        try:
            data = self.__stream__.next()
            self.__state__["token"] = self.resume_token
        except Exception as err:
            self.__state__["count"] += 1
            if time.time() - self.__state__["time"] > self.error_expiration:
                self.__state__["time"] = time.time()
                self.__state__["count"] = 1
            else:
                self.__state__["time"] = time.time()
                if self.__state__["count"] >= self.error_limit:
                    try:
                        self.__stream__.close()
                    except Exception:
                        pass
                    raise ValueError("More than {} errors have occured in {} seconds while watching for changes in {}".format(
                        self.error_limit, self.error_expiration, self.__watching_object__)) from err
            self.kwargs["resume_after"] = self.__state__["token"]
            self.__stream__ = self.__watching_object__.watch(self.pipeline, self.full_document, **self.kwargs)
            return self.__next__()

        return self._get_right_event(data)

    def try_next(self) -> typing.Any:
        """
        Try to get the next event without raising an exception and without waiting.

        Returns
        -------
        WatchEvent
        """
        try:
            data = self.__stream__.try_next()
            if data is not None:
                return self._get_right_event(data)
            return data
        except StopIteration:
            return None

    def close(self):
        """Closes the stream."""
        self.__closed__ = True
        self.__stream__.close()

    def resume(self):
        """Resume the stream from the last known state."""
        self.kwargs["resume_after"] = self.__state__["token"]
        self.__stream__ = self.__watching_object__.watch(self.pipeline, self.full_document, **self.kwargs)
        self.__closed__ = False

    @property
    def resume_token(self):
        """Get the resume token, which is used to resume the stream if failed."""
        return self.__stream__.resume_token

    @property
    def alive(self):
        """
        Get whether the stream is alive.

        Returns
        -------
        bool
            Whether the stream is alive.
        """
        return self._cursor.alive

    def __enter__(self):
        """
        Enter the context manager.

        Example
        -------
        >>> with db.watch() as stream: # <-- this line calls __enter__
        ...     for event in stream:
        ...         print(event)
        """
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Exit the context manager.

        Example
        -------
        >>> with db.watch() as stream:
        ...     for event in stream:
        ...         print(event)
        ... # <-- this line calls __exit__
        """
        self.close()

    def __iter__(self) -> typing.Iterator[WatchEvent]:
        """
        Returns the iterator.
        """
        return self

watch.py  Manages the watch cursor. An enum for the different event that can occur on MongoDB      Example     -------     >>> from yuno.direction import Operation     >>> Operation.INSERT     'insert' An object representing an event on MongoDB. The namespace the event occured in The id of the event Initialize the object with raw event data # insert, replace, delete, update # rename A db.watch(), db.collection.watch(), db.collection.object.watch() stream to get the different events. # the resume token # time of the last error # number of errors in the period # pipeline=pipeline, full_document=None, resume_after=resume_state["token"], max_await_time_ms=None, #    batch_size=None, collation=None, start_at_operation_time=None, session=None, start_after=None Initializes the stream.          Parameters         ----------         watching_object: pymongo.collection.Collection, pymongo.database.Database, pymongo.mongo_client.MongoClient             The object to watch.         pipeline: list[dict]             The pipeline to use.         full_document: bool, str             To return the full document when the event is an UpdateEvent.         error_limit: int             The number of errors before the stream is closed.         error_expiration: float             The number of seconds before the error count is reset.         kwargs:             The arguments to pass to the stream. # alias Get the next event (blocking operation)          Returns         -------         WatchEvent             The next event to occur on the object.          Example         -------         >>> document = db.collection.document         >>> watch_obj = document.watch()         >>> for event in watch_obj:         >>>     print(event) # only called when an event occurs on the database An internal function to get the right event type from raw event data.          Parameters         ----------         data: dict             The raw event data.          Returns         -------         WatchEvent             The event. Get the next event (blocking operation)          Returns         -------         WatchEvent             The next event to occur on the object.          Example         -------         >>> document = db.collection.document         >>> watch_obj = document.watch()         >>> for event in watch_obj:         >>>     print(event) # only called when an event occurs on the database Try to get the next event without raising an exception and without waiting.          Returns         -------         WatchEvent Closes the stream. Resume the stream from the last known state. Get the resume token, which is used to resume the stream if failed. Get whether the stream is alive.          Returns         -------         bool             Whether the stream is alive. Enter the context manager.          Example         -------         >>> with db.watch() as stream: # <-- this line calls __enter__         ...     for event in stream:         ...         print(event) Exit the context manager.          Example         -------         >>> with db.watch() as stream:         ...     for event in stream:         ...         print(event)         ... # <-- this line calls __exit__ Returns the iterator.