from datetime import datetime
from typing import (
    Any,
    Mapping,
    Sequence,
    Tuple,
    Type,
    TypeVar,
    Union,
    overload,
)
from uuid import UUID

from sqlalchemy import (
    Column,
)
from sqlalchemy.sql.elements import (
    SQLCoreOperations,
)
from sqlalchemy.sql.roles import TypedColumnsClauseRole

from ._expression_select_cls import Select, SelectOfScalar

_T = TypeVar("_T")


_TCCA = Union[
    TypedColumnsClauseRole[_T],
    SQLCoreOperations[_T],
    Type[_T],
]

# Generated TypeVars start


{% for i in range(number_of_types) %}
_TScalar_{{ i }} = TypeVar(
    "_TScalar_{{ i }}",
    Column,  # type: ignore
    Sequence,  # type: ignore
    Mapping,  # type: ignore
    UUID,
    datetime,
    float,
    int,
    bool,
    bytes,
    str,
    None,
)

_T{{ i }} = TypeVar("_T{{ i }}")

{% endfor %}

# Generated TypeVars end

@overload
def select(__ent0: _TCCA[_T0]) -> SelectOfScalar[_T0]: ...


@overload
def select(__ent0: _TScalar_0) -> SelectOfScalar[_TScalar_0]:  # type: ignore
    ...


# Generated overloads start

{% for signature in signatures %}

@overload
def select(  # type: ignore
    {% for arg in signature[0] %}{{ arg.name }}: {{ arg.annotation }}, {% endfor %}
    ) -> Select[Tuple[{%for ret in signature[1] %}{{ ret }} {% if not loop.last %}, {% endif %}{% endfor %}]]: ...

{% endfor %}

# Generated overloads end


def select(*entities: Any) -> Union[Select, SelectOfScalar]:  # type: ignore
    if len(entities) == 1:
        return SelectOfScalar(*entities)
    return Select(*entities)
