from typing import assert_type
from contextlib import AbstractAsyncContextManager
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql import Select

from fastsprout.core import Field
from fastsprout.core.fields.field_predicate import FieldPredicate
from fastsprout.data import DataR
from fastsprout.data.backend.sql import SQLEntity, SQLQuery, SQLSignpost
from fastsprout.data.backend.sql.capabilities import SQLAbilitable
from fastsprout.data.capabilities.protocols import Abilitable
from fastsprout.data.capabilities.query import BaseQuery
from fastsprout.data.entity import Entitieable
from fastsprout.data.streams.protocols import AsyncEntityStreamable, AsyncValueStreamable


class Hero(SQLEntity[int]):
    name: Field[str]


class Team(SQLEntity[int]):
    title: Field[str]


def queries(ability: SQLAbilitable[Hero], signpost: SQLSignpost) -> None:
    protocol: Abilitable[Hero, SQLQuery[Hero]] = ability
    assert_type(SQLAbilitable(signpost, entity=Hero), SQLAbilitable[Hero])
    assert_type(ability.query_for(), SQLQuery[Hero])
    assert_type(ability.query_for(Hero), SQLQuery[Hero])
    assert_type(protocol.query_for(), SQLQuery[Hero])
    assert_type(protocol.query_for(Hero), SQLQuery[Hero])
    assert_type(ability.query_for().where(Hero.name.eq("Batman")), SQLQuery[Hero])
    query: BaseQuery[Hero, Select[tuple[Hero]]] = ability.query_for()
    shared: Abilitable[Hero, BaseQuery[Hero, Select[tuple[Hero]]], SQLQuery[Hero]] = ability
    assert_type(shared.query_for(), SQLQuery[Hero])
    assert_type(shared.stream(query), AsyncEntityStreamable[Hero])
    assert_type(ability.stream(query), AsyncEntityStreamable[Hero])
    assert_type(ability.stream(query).sort(Hero.name), AsyncEntityStreamable[Hero])
    assert_type(ability.stream(query).to_values(Hero.name), AsyncValueStreamable[str])
    ability.query_for(Team)  # error: reportArgumentType; mypy: arg-type
    protocol.query_for(Team)  # error: reportArgumentType; mypy: arg-type


def routed(datar: DataR) -> None:
    ability = datar.ability(Hero)
    assert_type(ability.query_for().entity, type[Hero])
    assert_type(ability.query_for(Hero).entity, type[Hero])
    ability.query_for(Team)  # error: reportArgumentType; mypy: arg-type


def generic_identity[E: Entitieable[int, AbstractAsyncContextManager[AsyncSession]]](source: type[E], key: int) -> None:
    assert_type(source.id.eq(key), FieldPredicate[E])
