from collections.abc import AsyncIterator, Sequence
from contextlib import AbstractAsyncContextManager
from typing import assert_type
from uuid import UUID, uuid4
from sqlalchemy.orm import Mapped
from sqlalchemy.ext.asyncio import AsyncSession
from sqlmodel import Field as SQLField
from fastsprout.core import Field, ReadDTO, ReadField, WriteDTO
from fastsprout.data.backend.sql import SQLEntity, SQLQuery
from fastsprout.data.backend.sql.capabilities import SQLAbilitable
from fastsprout.data.capabilities.protocols import Abilitable
from fastsprout.data.entity import BaseEntity, Entitieable
from fastsprout.data.streams.protocols import AsyncEntityStreamable

class SQLUser(SQLEntity[int]):
    id: ReadField[int] = SQLField(primary_key=True)
    name: Field[str]

SQLUser(id=1, name="a")
assert_type(SQLUser(id=1, name="a").id, int)
assert_type(SQLUser.signpost().factory(), AbstractAsyncContextManager[AsyncSession])
assert_type(SQLUser.__signpost__.factory(), AbstractAsyncContextManager[AsyncSession])
SQLUser(id=None, name="a")  # error: reportArgumentType; mypy: arg-type
SQLUser(1, "a")  # error: reportCallIssue; mypy: call-arg

class InvalidSQLIdentity(SQLEntity[int | None]): ...  # error: reportInvalidTypeArguments; mypy: type-var
class InvalidIdentity(BaseEntity[int | None]): ...  # error: reportInvalidTypeArguments; mypy: type-var

def invalid_protocol(entity: Entitieable[int | None, None]) -> None: ...  # error: reportInvalidTypeArguments; mypy: type-var

def entities[T](source: type[T] | list[T]) -> list[T]:
    return []

assert_type(entities(SQLUser), list[SQLUser])

class Team(SQLEntity[UUID]):
    id: ReadField[UUID] = SQLField(default_factory=uuid4, primary_key=True)
    name: Field[str]

class TeamRead(ReadDTO, Team): ...
class TeamCreate(WriteDTO, Team): ...
class ChildCreate(TeamCreate): ...

assert_type(Team(name="team").id, UUID)
assert_type(Team.id.orm, Mapped[UUID])
assert_type(TeamRead(name="team").id, UUID)

def identity(entity: SQLEntity[UUID]) -> UUID:
    assert_type(entity.id, UUID)
    return entity.id

def forbidden(create: TeamCreate, child: ChildCreate) -> None:
    _ = create.id  # error: reportAttributeAccessIssue; mypy: call-overload
    create.id = uuid4()  # error: reportAttributeAccessIssue; mypy: arg-type
    identity_field = TeamCreate.id  # error: reportAttributeAccessIssue; mypy: type-var
    _ = child.id  # error: reportAttributeAccessIssue; mypy: call-overload


class UserQuery(SQLQuery[SQLUser]): ...


async def query_contract(ability: SQLAbilitable[SQLUser]) -> None:
    # given
    query = UserQuery(SQLUser)
    other = SQLQuery(Team)
    protocol: Abilitable[SQLUser, SQLQuery[SQLUser]] = ability

    # then
    assert_type(await ability.find_first(query), SQLUser | None)
    assert_type(await ability.find_exactly_one(query), SQLUser)
    assert_type(await ability.find_all(query), Sequence[SQLUser])
    assert_type(ability.stream(query), AsyncEntityStreamable[SQLUser])
    assert_type(ability.iter(query), AsyncIterator[SQLUser])
    assert_type(ability.iter_per(query), AsyncIterator[Sequence[SQLUser]])
    assert_type(await ability.delete_by_query(query), int)
    assert_type(await ability.update_by_query(query), int)
    assert_type(await protocol.find_first(query), SQLUser | None)

    await ability.find_first(other)  # error: reportArgumentType; mypy: arg-type
    await ability.find_exactly_one(other)  # error: reportArgumentType; mypy: arg-type
    await ability.find_all(other)  # error: reportArgumentType; mypy: arg-type
    ability.stream(other)  # error: reportArgumentType; mypy: arg-type
    ability.iter(other)  # error: reportArgumentType; mypy: arg-type
    ability.iter_per(other)  # error: reportArgumentType; mypy: arg-type
    await ability.delete_by_query(other)  # error: reportArgumentType; mypy: arg-type
    await ability.update_by_query(other)  # error: reportArgumentType; mypy: arg-type
    await protocol.find_first(other)  # error: reportArgumentType; mypy: arg-type
