from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.user import User
from app.repositories.protocols import IUserRepository


class UserRepository(IUserRepository):
    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def get_by_telegram_id(self, telegram_id: int) -> User | None:
        result = await self._session.execute(
            select(User).where(User.telegram_id == telegram_id).limit(1)
        )
        return result.scalar_one_or_none()

    async def create(
        self,
        *,
        telegram_id: int,
        username: str | None,
        first_name: str | None,
        referrer_id: int | None,
    ) -> User:
        user = User(
            telegram_id=telegram_id,
            username=username,
            first_name=first_name,
            referrer_id=referrer_id,
        )
        self._session.add(user)
        await self._session.flush()
        await self._session.refresh(user)
        return user

    async def update_profile(
        self,
        user: User,
        *,
        username: str | None,
        first_name: str | None,
    ) -> User:
        user.username = username
        user.first_name = first_name
        await self._session.flush()
        await self._session.refresh(user)
        return user

    async def set_referrer_if_unset(
        self,
        user: User,
        referrer_internal_id: int,
    ) -> None:
        if user.referrer_id is not None:
            return
        if user.id == referrer_internal_id:
            return
        user.referrer_id = referrer_internal_id
        await self._session.flush()