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

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


class UserService:
    def __init__(self, user_repository: UserRepository) -> None:
        self.user_repository = user_repository

    async def add_user(
        self, session: AsyncSession, user_id: int, referal: Optional[int] = None
    ) -> User:
        user = await self.user_repository.get_user_by_id(session, user_id)
        if user:
            return user

        referal = int(referal) if str(referal).isdigit() else None
        model = NewUserDTO(id=user_id, referal=referal)
        return await self.user_repository.create_user(session, model)