malenia/services/user_service.py

import datetime
import logging
import math

from aiogram.types import Message, User as TelegramUser
from remnawave import RemnawaveSDK
from sqlalchemy.ext.asyncio import AsyncSession

from bot.texts import NO_SUBSCRIPTION, NOTHING_PLACEHOLDER, SUBSCRIPTION_INFORMATION
from db.models.users import User
from misc import utils
from repositories import UserRepository
from repositories.subscriptions import SubscriptionRepository
from schemas.subscriptions import InternalSquads, SubscriptionPlan, SubscriptionStatus
from schemas.users import NewUserDTO
from services import rw
from config import settings

logger = logging.getLogger(__name__)


class UserServiceError(Exception): ...


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

    @staticmethod
    def format_subscription_message(
        rw_user: rw.RWUserInfo | None, link: str | None = None, autorenew: bool = False
    ):
        if rw_user is None:
            # Пользователь не найден в Remnawave
            return NO_SUBSCRIPTION

        status_icon, status_text = utils.format_status(rw_user.status)
        expire_date = rw_user.expire_at.strftime("%d.%m.%Y")
        days_left = utils.format_days_left(rw_user.expire_at)
        used_traffic = utils.format_traffic(rw_user.used_traffic_bytes, rw_user.traffic_limit_bytes)
        hwid_limit = utils.format_hwid_limit(rw_user.hwid_device_limit)
        link = link or NOTHING_PLACEHOLDER
        autorenew_text = "🟢" if autorenew else "🔴"

        return SUBSCRIPTION_INFORMATION.format(
            status_icon=status_icon,
            status=status_text,
            expire_date=expire_date,
            days_left=days_left,
            used_traffic=used_traffic,
            hwid_limit=hwid_limit,
            link=link,
            autorenew=autorenew_text,
        )

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

        rw_user = await rw.get_user_by_telegram_id(rw_sdk, user_id)
        link = utils.get_subscription_link(rw_user.short_uuid) if rw_user else None
        referal = int(referal) if str(referal).isdigit() and user_id != int(referal) else None

        model = NewUserDTO(id=user_id, referal=referal, link=link)
        return await self.user_repository.create_user(session, model)

    async def update_user_link(
        self, session: AsyncSession, user_id: int, rw_sdk: RemnawaveSDK
    ) -> User | None:
        rw_user = await rw.get_user_by_telegram_id(rw_sdk, user_id)

        user = await self.user_repository.get_user_by_id(session, user_id)

        link = utils.get_subscription_link(rw_user.short_uuid) if rw_user else None
        user = await self.user_repository.update_user_link(session, user_id, link)

        return user

    async def send_promotion_to_users(self, session: AsyncSession, msg: Message) -> list[User]:
        users = await self.user_repository.get_users(session)

        sent: list[User] = []
        for user in users:
            try:
                await msg.copy_to(user.id)
                sent.append(user)
            except Exception:
                logger.critical("failed to send a message to user %d.", user.id, exc_info=True)

        return sent

    async def buy_subscription(
        self,
        session: AsyncSession,
        subscription: SubscriptionPlan,
        rw_sdk: RemnawaveSDK,
        user: TelegramUser,
    ):
        duration = utils.convert_duration_to_timedelta(subscription.duration)
        squads = [InternalSquads.DEFAULT]
        if subscription.whitelists:
            squads.append(InternalSquads.WHITELISTS)

        ## If subscription exists, add.
        rw_user = await rw.get_user_by_telegram_id(rw_sdk, user.id)
        if rw_user:
            db_subscription = await self.subscription_repository.get_subscription_by_user_id(
                session, user.id
            )
            if db_subscription:
                db_subscription = await self.subscription_repository.update_sub(
                    session,
                    user.id,
                    end_date=rw_user.expire_at + duration,
                    autorenew=False,
                    status=SubscriptionStatus.ACTIVE,
                    devices=int(subscription.devices),
                    whitelists=subscription.whitelists,
                    plan_price_paid=utils.calculate_price(subscription),
                    duration_days=duration.days,
                )
            else:
                db_subscription = await self.subscription_repository.create_subscription(
                    session,
                    user_id=user.id,
                    end_date=rw_user.expire_at + duration,
                    autorenew=False,
                    status=SubscriptionStatus.ACTIVE,
                    devices=int(subscription.devices),
                    whitelists=subscription.whitelists,
                    plan_price_paid=utils.calculate_price(subscription),
                    duration_days=duration.days,
                )

            if not utils.is_same_plan(rw_user, subscription):
                new_expire = datetime.datetime.now(datetime.UTC) + duration
                delta_days = (new_expire - rw_user.expire_at).days
                await rw.add_days(rw_sdk, rw_user.uuid, delta_days)
            else:
                await rw.add_days(rw_sdk, rw_user.uuid, duration.days)

            await rw.update_squads(rw_sdk, rw_user.uuid, squads)
            await rw.set_hwid_limit(rw_sdk, rw_user.uuid, subscription.devices)

            await self.update_user_link(session, user.id, rw_sdk)
            return db_subscription

        # Create
        expire = datetime.datetime.now(datetime.UTC) + duration
        username = "tg_" + str(user.username or user.id)

        rw_user = await rw.create_user(
            rw_sdk,
            username=username,
            telegram_id=user.id,
            expire_at=expire,
            hwid_device_limit=subscription.devices,
            squad_uuids=squads,
        )

        db_subscription = await self.subscription_repository.get_subscription_by_user_id(
            session, user.id
        )
        if db_subscription:
            db_subscription = await self.subscription_repository.update_sub(
                session,
                user.id,
                end_date=expire,
                autorenew=False,
                status=SubscriptionStatus.ACTIVE,
                devices=int(subscription.devices),
                whitelists=subscription.whitelists,
                plan_price_paid=utils.calculate_price(subscription),
                duration_days=duration.days,
            )
        else:
            db_subscription = await self.subscription_repository.create_subscription(
                session,
                user_id=user.id,
                end_date=expire,
                autorenew=False,
                status=SubscriptionStatus.ACTIVE,
                devices=int(subscription.devices),
                whitelists=subscription.whitelists,
                plan_price_paid=utils.calculate_price(subscription),
                duration_days=duration.days,
            )

        await self.update_user_link(session, user.id, rw_sdk)

        return db_subscription

    async def calculate_plan_credit(
        self,
        session: AsyncSession,
        user_id: int,
        new_plan: SubscriptionPlan,
        rw_sdk: RemnawaveSDK,
    ) -> int:
        rw_user = await rw.get_user_by_telegram_id(rw_sdk, user_id)

        if not rw_user or rw_user.expire_at <= datetime.datetime.now(datetime.UTC):
            return 0

        if utils.is_same_plan(rw_user, new_plan):
            return 0

        db_sub = await self.subscription_repository.get_subscription_by_user_id(session, user_id)
        if not db_sub or not db_sub.plan_price_paid or not db_sub.duration_days:
            return 0

        remaining_days = max(0, (rw_user.expire_at - datetime.datetime.now(datetime.UTC)).days)
        return math.floor(db_sub.plan_price_paid * remaining_days / db_sub.duration_days)

    async def launch_promotion(
        self,
        session: AsyncSession,
        promo_msg: Message
    ) -> int:
        users = await self.user_repository.get_users(session)

        sent = 0
        for user in users:
            if user.id in settings.admins:
                continue
            try:
                await promo_msg.copy_to(user.id)
                sent += 1
            except Exception:
                logger.exception("caught an exception while launching promo, skipping user...")
        
        return sent