"""
Transaction management for the async neomodel module.
"""

import warnings
from functools import wraps
from inspect import iscoroutinefunction
from typing import Any, Callable

from neo4j.api import Bookmarks
from neo4j.exceptions import ClientError

from neomodel._async_compat.util import AsyncUtil
from neomodel.async_.database import AsyncDatabase
from neomodel.constants import NOT_COROUTINE_ERROR
from neomodel.exceptions import UniqueProperty


class AsyncTransactionProxy:
    def __init__(
        self,
        db: AsyncDatabase,
        access_mode: str | None = None,
        parallel_runtime: bool | None = False,
    ):
        self.db: AsyncDatabase = db
        self.access_mode: str | None = access_mode
        self.parallel_runtime: bool | None = parallel_runtime
        self.bookmarks: Bookmarks | None = None
        self.last_bookmarks: Bookmarks | None = None

    async def __aenter__(self) -> "AsyncTransactionProxy":
        if self.parallel_runtime and not await self.db.parallel_runtime_available():
            warnings.warn(
                "Parallel runtime is only available in Neo4j Enterprise Edition 5.13 and above. "
                "Reverting to default runtime.",
                UserWarning,
            )
            self.parallel_runtime = False
        self.db._parallel_runtime = self.parallel_runtime
        await self.db.begin(access_mode=self.access_mode, bookmarks=self.bookmarks)
        self.bookmarks = None
        return self

    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        self.db._parallel_runtime = False
        if exc_value:
            await self.db.rollback()

        if (
            exc_type is ClientError
            and exc_value.code == "Neo.ClientError.Schema.ConstraintValidationFailed"
        ):
            raise UniqueProperty(exc_value.message)

        if not exc_value:
            self.last_bookmarks = await self.db.commit()

    def __call__(self, func: Callable) -> Callable:
        if AsyncUtil.is_async_code and not iscoroutinefunction(func):
            raise TypeError(NOT_COROUTINE_ERROR)

        @wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Callable:
            async with self:
                return await func(*args, **kwargs)

        return wrapper

    @property
    def with_bookmark(self) -> "BookmarkingAsyncTransactionProxy":
        return BookmarkingAsyncTransactionProxy(self.db, self.access_mode)


class BookmarkingAsyncTransactionProxy(AsyncTransactionProxy):
    def __call__(self, func: Callable) -> Callable:
        if AsyncUtil.is_async_code and not iscoroutinefunction(func):
            raise TypeError(NOT_COROUTINE_ERROR)

        async def wrapper(*args: Any, **kwargs: Any) -> tuple[Any, None]:
            self.bookmarks = kwargs.pop("bookmarks", None)

            async with self:
                result = await func(*args, **kwargs)
                self.last_bookmarks = None

            return result, self.last_bookmarks

        return wrapper


class ImpersonationHandler:
    def __init__(self, db: AsyncDatabase, impersonated_user: str):
        self.db = db
        self.impersonated_user = impersonated_user

    def __enter__(self) -> "ImpersonationHandler":
        self.db.impersonated_user = self.impersonated_user
        return self

    def __exit__(
        self, exception_type: Any, exception_value: Any, exception_traceback: Any
    ) -> None:
        self.db.impersonated_user = None

        print("\nException type:", exception_type)
        print("\nException value:", exception_value)
        print("\nTraceback:", exception_traceback)

    def __call__(self, func: Callable) -> Callable:
        def wrapper(*args: Any, **kwargs: Any) -> Callable:
            with self:
                return func(*args, **kwargs)

        return wrapper
