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

from db.models import User
from dto.users import NewUserDTO


class UserRepository:
    async def get_user_by_id(
        self, session: AsyncSession, user_id: int
    ) -> Optional[User]:
        stmt = select(User).where(User.id == user_id)
        result = await session.execute(stmt)

        return result.scalar_one_or_none()

    async def create_user(self, session: AsyncSession, user_data: NewUserDTO) -> User:
        user = User(id=user_data.id, referal_id=user_data.referal)
        session.add(user)
        await session.commit()

        return user

    async def update_user_topic(
        self, session: AsyncSession, user_id: int, thread_id: int
    ) -> Optional[User]:
        user = await self.get_user_by_id(session, user_id=user_id)
        if not user:
            return None
        user.support_thread_id = thread_id

        await session.commit()
        return user