from sqlmodel import select
from sqlmodel.ext.asyncio.session import AsyncSession

from utils.db.models import User


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

    async def get(self, user_id: int) -> User | None:
        return await self._session.get(User, user_id)

    async def by_tg_id(self, tg_id: int) -> User | None:
        stmt = select(User).where(User.tg_id == tg_id)
        return (await self._session.exec(stmt)).first()

    async def upsert(
        self,
        tg_id: int,
        *,
        username: str | None = None,
        first_name: str | None = None,
        last_name: str | None = None,
    ) -> User:
        user = await self.by_tg_id(tg_id)
        if user is None:
            user = User(tg_id=tg_id)
        user.username = username
        user.first_name = first_name
        user.last_name = last_name
        self._session.add(user)
        await self._session.commit()
        await self._session.refresh(user)
        return user