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

from sqlalchemy import Column
from sqlalchemy.orm import InstrumentedAttribute
from sqlalchemy.sql.elements import ColumnClause
from sqlalchemy.sql.expression import Select as _Select

_TSelect = TypeVar("_TSelect")

# Workaround Generics incompatibility in Python 3.6
# Ref: https://github.com/python/typing/issues/449#issuecomment-316061322
if sys.version_info.minor >= 7:

    class Select(_Select, Generic[_TSelect]):
        inherit_cache = True

    # This is not comparable to sqlalchemy.sql.selectable.ScalarSelect, that has a different
    # purpose. This is the same as a normal SQLAlchemy Select class where there's only one
    # entity, so the result will be converted to a scalar by default. This way writing
    # for loops on the results will feel natural.
    class SelectOfScalar(_Select, Generic[_TSelect]):
        inherit_cache = True

else:
    from typing import GenericMeta  # type: ignore

    class GenericSelectMeta(GenericMeta, _Select.__class__):  # type: ignore
        pass

    class _Py36Select(_Select, Generic[_TSelect], metaclass=GenericSelectMeta):
        inherit_cache = True

    class _Py36SelectOfScalar(_Select, Generic[_TSelect], metaclass=GenericSelectMeta):
        inherit_cache = True

    # Cast them for editors to work correctly, from several tricks tried, this works
    # for both VS Code and PyCharm
    Select = cast("Select", _Py36Select)  # type: ignore
    SelectOfScalar = cast("SelectOfScalar", _Py36SelectOfScalar)  # type: ignore


if TYPE_CHECKING:  # pragma: no cover
    from ..main import SQLModel

# 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,
)

_TModel_{{ i }} = TypeVar("_TModel_{{ i }}", bound="SQLModel")

{% endfor %}

# Generated TypeVars end

@overload
def select(entity_0: _TScalar_0, **kw: Any) -> SelectOfScalar[_TScalar_0]:  # type: ignore
    ...


@overload
def select(entity_0: Type[_TModel_0], **kw: Any) -> SelectOfScalar[_TModel_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 %}**kw: Any,
    ) -> Select[Tuple[{%for ret in signature[1] %}{{ ret }} {% if not loop.last %}, {% endif %}{% endfor %}]]:
    ...

{% endfor %}

# Generated overloads end

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


# TODO: add several @overload from Python types to SQLAlchemy equivalents
def col(column_expression: Any) -> ColumnClause:  # type: ignore
    if not isinstance(column_expression, (ColumnClause, Column, InstrumentedAttribute)):
        raise RuntimeError(f"Not a SQLAlchemy column: {column_expression}")
    return column_expression