from datetime import datetime
from typing import (
    TYPE_CHECKING,
    Any,
    Generic,
    Mapping,
    Sequence,
    Tuple,
    Type,
    TypeVar,
    Union,
    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")

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

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