from typing import Any
from typing import Callable
from typing import ContextManager
from typing import Iterable
from typing import Iterator
from typing import Mapping
from typing import Optional
from typing import Sequence
from typing import TypeVar
from typing import Union

from .base import StartableContext
from .engine import AsyncConnection
from .engine import AsyncEngine
from .result import AsyncResult
from ... import util
from ...engine import Result
from ...engine.base import _ExecutionOptions
from ...orm import Session
from ...orm.session import _IdentityMap
from ...sql import ClauseElement
from ...sql import Executable

_T = TypeVar("_T")
_TAsyncSession = TypeVar("_TAsyncSession", bound=AsyncSession)
_TAsyncSessionTransaction = TypeVar(
    "_TAsyncSessionTransaction", bound=AsyncSessionTransaction
)

class AsyncSession:
    dispatch: Any = ...
    bind: Any = ...
    binds: Any = ...
    sync_session: Session = ...
    def __init__(
        self,
        bind: Optional[Union[AsyncConnection, AsyncEngine]] = ...,
        binds: Optional[
            Mapping[object, Union[AsyncConnection, AsyncEngine]]
        ] = ...,
        **kw: Any,
    ) -> None: ...
    async def refresh(
        self,
        instance: Any,
        attribute_names: Optional[Any] = ...,
        with_for_update: Optional[Any] = ...,
    ) -> None: ...
    async def run_sync(
        self, fn: Callable[..., _T], *arg: Any, **kw: Any
    ) -> _T: ...
    async def execute(
        self,
        statement: Executable,
        params: Optional[
            Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]
        ] = ...,
        execution_options: Optional[_ExecutionOptions] = ...,
        bind_arguments: Optional[Mapping[str, Any]] = ...,
        **kw: Any,
    ) -> Result: ...
    async def scalar(
        self,
        statement: Executable,
        params: Optional[
            Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]
        ] = ...,
        execution_options: _ExecutionOptions = ...,
        bind_arguments: Optional[Mapping[str, Any]] = ...,
        **kw: Any,
    ) -> Any: ...
    async def get(
        self,
        entity: Any,
        ident: Any,
        options: Optional[Any] = ...,
        populate_existing: bool = ...,
        with_for_update: Optional[Any] = ...,
        identity_token: Optional[Any] = ...,
    ) -> Any: ...
    async def stream(
        self,
        statement: Any,
        params: Optional[
            Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]
        ] = ...,
        execution_options: _ExecutionOptions = ...,
        bind_arguments: Optional[Mapping[str, Any]] = ...,
        **kw: Any,
    ) -> AsyncResult: ...
    async def delete(self, instance: Any) -> None: ...
    async def merge(self, instance: _T, load: bool = ...) -> _T: ...
    async def flush(self, objects: Optional[Any] = ...) -> None: ...
    async def connection(self) -> AsyncConnection: ...
    def begin(self, **kw: Any) -> AsyncSessionTransaction: ...
    def begin_nested(self, **kw: Any) -> AsyncSessionTransaction: ...
    async def rollback(self) -> None: ...
    async def commit(self) -> None: ...
    async def close(self) -> None: ...
    @classmethod
    async def close_all(self) -> None: ...
    async def __aenter__(self: _TAsyncSession) -> _TAsyncSession: ...
    async def __aexit__(
        self, type_: Any, value: Any, traceback: Any
    ) -> None: ...
    # copied via create_proxy_methods
    def __contains__(self, instance: Any) -> bool: ...
    def __iter__(self) -> Iterator[Any]: ...
    def add(self, instance: Any, _warn: bool = ...) -> None: ...
    def add_all(self, instances: Any) -> None: ...
    def expire(
        self, instance: Any, attribute_names: Optional[Iterable[str]] = ...
    ) -> None: ...
    def expire_all(self) -> None: ...
    def expunge(self, instance: Any) -> None: ...
    def expunge_all(self) -> None: ...
    def get_bind(
        self,
        mapper: Optional[Any] = ...,
        clause: Optional[ClauseElement] = ...,
        bind: Optional[Union[AsyncConnection, AsyncEngine]] = ...,
        _sa_skip_events: Optional[Any] = ...,
        _sa_skip_for_implicit_returning: bool = ...,
    ) -> Union[AsyncConnection, AsyncEngine]: ...
    def is_modified(
        self, instance: Any, include_collections: bool = ...
    ) -> bool: ...
    def in_transaction(self) -> bool: ...
    @property
    def dirty(self) -> util.IdentitySet[Any]: ...
    @property
    def deleted(self) -> util.IdentitySet[Any]: ...
    @property
    def new(self) -> util.IdentitySet[Any]: ...
    identity_map: _IdentityMap
    @property
    def is_active(self) -> bool: ...
    autoflush: bool
    @property
    def no_autoflush(
        self: _TAsyncSession,
    ) -> ContextManager[_TAsyncSession]: ...
    @util.memoized_property
    def info(self) -> Mapping[Any, Any]: ...

class _AsyncSessionContextManager:
    async_session: AsyncSession = ...
    trans: AsyncSessionTransaction = ...
    def __init__(self, async_session: AsyncSession) -> None: ...
    async def __aenter__(self) -> AsyncSession: ...
    async def __aexit__(
        self, type_: Any, value: Any, traceback: Any
    ) -> None: ...

class AsyncSessionTransaction(StartableContext):
    session: AsyncSession = ...
    nested: bool = ...
    sync_transaction: Optional[Any] = ...
    def __init__(self, session: AsyncSession, nested: bool = ...) -> None: ...
    @property
    def is_active(self) -> bool: ...
    async def rollback(self) -> Optional[AsyncSessionTransaction]: ...
    async def commit(self) -> Optional[AsyncSessionTransaction]: ...
    async def start(
        self: _TAsyncSessionTransaction,
    ) -> _TAsyncSessionTransaction: ...
    async def __aexit__(
        self, type_: Any, value: Any, traceback: Any
    ) -> None: ...
