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

from db.models import User
from db.models.transactions import BalanceTransaction, BalanceTxType


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

    async def get_user_by_id(self, id: int) -> User | None:
        stmt = select(User).where(User.id == id)
        res = await self.session.execute(stmt)
        return res.scalar_one_or_none()

    async def get_user_by_telegram_id(self, telegram_id: int) -> User | None:
        stmt = select(User).where(User.telegram_id == telegram_id)
        res = await self.session.execute(stmt)
        return res.scalar_one_or_none()

    async def get_user_by_username(self, username: str) -> User | None:
        stmt = select(User).where(User.username == username)
        res = await self.session.execute(stmt)
        return res.scalar_one_or_none()

    async def get_user_by_ref_code(self, ref_code: str) -> User | None:
        stmt = select(User).where(User.referal_code == ref_code)
        res = await self.session.execute(stmt)
        return res.scalar_one_or_none()

    async def create(
        self,
        *,
        username: str | None = None,
        hashed_password: str | None = None,
        telegram_id: int | None = None,
        referal_id: int | None = None,
    ) -> User:
        obj = User(
            username=username,
            hashed_password=hashed_password,
            telegram_id=telegram_id,
            referal_id=referal_id,
        )
        self.session.add(obj)
        await self.session.commit()

        return obj

    async def increase_balance(
        self, user_id: int, amount: float, tx_type: BalanceTxType, description: str
    ):
        user = await self.get_user_by_id(user_id)

        if not user:
            return

        obj = BalanceTransaction(
            user_id=user_id,
            amount=amount,
            tx_type=tx_type,
            balance_before=user.balance,
            balance_after=user.balance + amount,
            description=description,
        )
        self.session.add(obj)

        user.balance += amount
        await self.session.commit()
        return user

    async def update_telegram_id(self, user: User, telegram_id: int) -> User:
        user.telegram_id = telegram_id
        await self.session.commit()

        return user