from __future__ import annotations

from typing import (
    Mapping, 
    Sequence, 
    TypeVar, 
    TYPE_CHECKING, 
    Union, 
    Tuple, 
    Callable, 
    Any, 
    List, 
    Protocol,
    Optional, 
    AsyncIterator, 
    KeysView
)
from sqlalchemy.engine.cursor import ResultProxy
from sqlalchemy.ext.asyncio import AsyncResult
from sqlalchemy.sql.visitors import Visitable

if TYPE_CHECKING:
    from .schemas import Schema
    from .engine import Engine, Connection
    from .results import Row
    from subway.models import Model

AnyRow = Union['Row', Tuple[Any, ...]]
UniqueStrategy = Callable[..., Any]
Bind = Union['Engine', 'Connection']
MappingValue = Mapping[str, Any]

ResultT = TypeVar('ResultT', bound='SupportsRow')
SchemaT = TypeVar("SchemaT", bound="Schema")
SupportsWhereT = TypeVar('SupportsWhereT', bound="SupportsWhere")

Entity = Union[SchemaT, 'Model', Mapping[str, Any], Sequence[Any]]

class SupportsRow(Protocol):
    @classmethod
    def from_row(cls, row: Row) -> SupportsRow:
        ...

class IteratorResult(ResultProxy):
    pass

class FrozenResult(Protocol):
    def __init__(self, result: AsyncResult) -> None: ...
    def rewrite_rows(self) -> List[AnyRow]: ...
    def with_new_rows(self, rows: List[AnyRow]) -> FrozenResult: ...
    def __call__(self) -> IteratorResult: ...

class MappingResult(Protocol):
    def keys(self) -> KeysView[str]: ...
    def unique(self, strategy: Optional[UniqueStrategy]=None) -> MappingResult: ...
    async def fetchone(self) -> Optional[MappingValue]: ...
    async def fetchmany(self, size: Optional[int]=None) -> List[MappingValue]: ...
    async def fetchall(self) -> List[MappingValue]: ...
    async def all(self) -> List[MappingValue]: ...
    async def first(self) -> Optional[MappingValue]: ...
    async def one(self) -> MappingValue: ...
    async def one_or_none(self) -> Optional[MappingValue]: ...
    async def partitions(self, size: Optional[int]=None) -> AsyncIterator[MappingValue]: ...
    async def __aiter__(self) -> AsyncIterator[MappingValue]: ...

class ScalarsResult(Protocol):
    def unique(self, strategy: Optional[UniqueStrategy]=None) -> ScalarsResult: ...
    async def fetchone(self) -> Optional[Any]: ...
    async def fetchmany(self, size: Optional[int]=None) -> List[Any]: ...
    async def fetchall(self) -> List[Any]: ...
    async def all(self) -> List[Any]: ...
    async def first(self) -> Optional[Any]: ...
    async def one(self) -> Any: ...
    async def one_or_none(self) -> Optional[Any]: ...
    async def partitions(self, size: Optional[int]=None) -> AsyncIterator[Any]: ...
    async def __aiter__(self) -> AsyncIterator[Any]: ...

class SupportsWhere(Protocol):
    def where(self, whereclause: Union[str, bool, Visitable]) -> SupportsWhere: ...

from __future__ import annotations

from typing import (
    Mapping, 
    Sequence, 
    TypeVar, 
    TYPE_CHECKING, 
    Union, 
    Tuple, 
    Callable, 
    Any, 
    List, 
    Protocol,
    Optional, 
    AsyncIterator, 
    KeysView
)
from sqlalchemy.engine.cursor import ResultProxy
from sqlalchemy.ext.asyncio import AsyncResult
from sqlalchemy.sql.visitors import Visitable

if TYPE_CHECKING:
    from .schemas import Schema
    from .engine import Engine, Connection
    from .results import Row
    from subway.models import Model

AnyRow = Union['Row', Tuple[Any, ...]]
UniqueStrategy = Callable[..., Any]
Bind = Union['Engine', 'Connection']
MappingValue = Mapping[str, Any]

ResultT = TypeVar('ResultT', bound='SupportsRow')
SchemaT = TypeVar("SchemaT", bound="Schema")
SupportsWhereT = TypeVar('SupportsWhereT', bound="SupportsWhere")

Entity = Union[SchemaT, 'Model', Mapping[str, Any], Sequence[Any]]

class SupportsRow(Protocol):
    @classmethod
    def from_row(cls, row: Row) -> SupportsRow:
        ...

class IteratorResult(ResultProxy):
    pass

class FrozenResult(Protocol):
    def __init__(self, result: AsyncResult) -> None: ...
    def rewrite_rows(self) -> List[AnyRow]: ...
    def with_new_rows(self, rows: List[AnyRow]) -> FrozenResult: ...
    def __call__(self) -> IteratorResult: ...

class MappingResult(Protocol):
    def keys(self) -> KeysView[str]: ...
    def unique(self, strategy: Optional[UniqueStrategy]=None) -> MappingResult: ...
    async def fetchone(self) -> Optional[MappingValue]: ...
    async def fetchmany(self, size: Optional[int]=None) -> List[MappingValue]: ...
    async def fetchall(self) -> List[MappingValue]: ...
    async def all(self) -> List[MappingValue]: ...
    async def first(self) -> Optional[MappingValue]: ...
    async def one(self) -> MappingValue: ...
    async def one_or_none(self) -> Optional[MappingValue]: ...
    async def partitions(self, size: Optional[int]=None) -> AsyncIterator[MappingValue]: ...
    async def __aiter__(self) -> AsyncIterator[MappingValue]: ...

class ScalarsResult(Protocol):
    def unique(self, strategy: Optional[UniqueStrategy]=None) -> ScalarsResult: ...
    async def fetchone(self) -> Optional[Any]: ...
    async def fetchmany(self, size: Optional[int]=None) -> List[Any]: ...
    async def fetchall(self) -> List[Any]: ...
    async def all(self) -> List[Any]: ...
    async def first(self) -> Optional[Any]: ...
    async def one(self) -> Any: ...
    async def one_or_none(self) -> Optional[Any]: ...
    async def partitions(self, size: Optional[int]=None) -> AsyncIterator[Any]: ...
    async def __aiter__(self) -> AsyncIterator[Any]: ...

class SupportsWhere(Protocol):
    def where(self, whereclause: Union[str, bool, Visitable]) -> SupportsWhere: ...
