malenia/services/renewal_service.py

import datetime
import logging

from aiogram import Bot
from remnawave import RemnawaveSDK
from sqlalchemy.ext.asyncio import AsyncSession

from bot.texts import AUTORENEWAL_INSUFFICIENT_FUNDS, AUTORENEWAL_SUCCESS
from config import settings
from db.models.subscriptions import Subscription
from db.models.transactions import BalanceTxType
from misc import utils
from repositories.subscriptions import SubscriptionRepository
from repositories.users import UserRepository
from schemas.subscriptions import InternalSquads, SubscriptionPlan, SubscriptionStatus
from services import rw as rw_service
from services.admin_notifications import notify_subscription_renewal

logger = logging.getLogger(__name__)


class RenewalServiceError(Exception):
    """Base exception for renewal service errors."""


class InsufficientBalanceError(RenewalServiceError):
    """Raised when user has insufficient balance for renewal."""


class RenewalService:
    def __init__(
        self,
        subscription_repository: SubscriptionRepository,
        user_repository: UserRepository,
        rw_sdk: RemnawaveSDK,
        bot: Bot,
    ) -> None:
        self.subscription_repository = subscription_repository
        self.user_repository = user_repository
        self.rw_sdk = rw_sdk
        self.bot = bot

    async def _calculate_renewal_price(self, subscription: Subscription) -> int:
        duration = utils.convert_int_to_duration(subscription.duration_days)
        plan = SubscriptionPlan(
            devices=subscription.devices,
            duration=duration,
            whitelists=subscription.whitelists,
        )
        return utils.calculate_price(plan)

    async def _renew_subscription_in_rw(self, subscription: Subscription) -> None:
        rw_user = await rw_service.get_user_by_telegram_id(self.rw_sdk, subscription.user_id)
        if not rw_user:
            raise RenewalServiceError(f"RW user not found for user_id={subscription.user_id}")

        squads = [InternalSquads.DEFAULT]
        if subscription.whitelists:
            squads.append(InternalSquads.WHITELISTS)

        await rw_service.add_days(self.rw_sdk, rw_user.uuid, subscription.duration_days)
        await rw_service.update_squads(self.rw_sdk, rw_user.uuid, squads)
        await rw_service.set_hwid_limit(self.rw_sdk, rw_user.uuid, subscription.devices)

    async def _notify_user(self, user_id: int, message: str) -> None:
        try:
            await self.bot.send_message(user_id, message)
        except Exception:
            logger.warning("Failed to notify user %d about renewal", user_id, exc_info=True)

    async def _process_single_renewal(
        self, session: AsyncSession, subscription: Subscription
    ) -> bool:
        user_id = subscription.user_id
        price = await self._calculate_renewal_price(subscription)

        user = await self.user_repository.get_user_by_id(session, user_id)
        if not user:
            logger.error("User %d not found for autorenewal", user_id)
            return False

        if user.balance < price:
            logger.info(
                "Insufficient balance for user %d: balance=%d, price=%d",
                user_id,
                user.balance,
                price,
            )
            await self.subscription_repository.update_subscription_autorenew(
                session, user_id, False
            )
            await self._notify_user(user_id, AUTORENEWAL_INSUFFICIENT_FUNDS)
            return False

        renewed_user = await self.user_repository.decrease_user_balance(
            session,
            user_id,
            price,
            tx_type=BalanceTxType.AUTORENEW,
            description=f"autorenewal: devices={subscription.devices}, duration={subscription.duration_days}d, whitelists={subscription.whitelists}",
        )

        if not renewed_user:
            logger.error("Failed to decrease balance for user %d", user_id)
            return False

        try:
            await self._renew_subscription_in_rw(subscription)
        except RenewalServiceError:
            logger.exception("Failed to renew subscription in RW for user %d", user_id)
            await self.user_repository.increase_user_balance(
                session,
                user_id,
                price,
                tx_type=BalanceTxType.REFUND,
                description=f"refund for failed autorenewal: devices={subscription.devices}, duration={subscription.duration_days}d",
            )
            return False

        new_end_date = subscription.end_date + datetime.timedelta(days=subscription.duration_days)
        await self.subscription_repository.update_sub(
            session,
            user_id,
            end_date=new_end_date,
            autorenew=True,
            status=SubscriptionStatus.ACTIVE,
            devices=subscription.devices,
            whitelists=subscription.whitelists,
            duration_days=subscription.duration_days,
            plan_price_paid=price,
        )

        await self._notify_user(
            user_id,
            AUTORENEWAL_SUCCESS.format(
                price=price,
                end_date=new_end_date.strftime("%d.%m.%Y"),
            ),
        )
        await notify_subscription_renewal(
            self.bot,
            user_id=user_id,
            amount=price,
            balance_before=renewed_user.balance + price,
            balance_after=renewed_user.balance,
            devices=subscription.devices,
            duration_days=subscription.duration_days,
            whitelists=subscription.whitelists,
            end_date=new_end_date.strftime("%d.%m.%Y"),
        )

        logger.info("Successfully renewed subscription for user %d", user_id)
        return True

    async def process_autorenewals(self, session: AsyncSession) -> dict:
        subscriptions = await self.subscription_repository.get_subscriptions_for_autorenewal(
            session, days_before_expiry=settings.autorenew_days_before
        )

        logger.info("Found %d subscriptions for autorenewal", len(subscriptions))

        results = {"success": 0, "failed": 0, "insufficient_balance": 0}

        for subscription in subscriptions:
            try:
                success = await self._process_single_renewal(session, subscription)
                if success:
                    results["success"] += 1
                else:
                    user = await self.user_repository.get_user_by_id(session, subscription.user_id)
                    if user and user.balance < await self._calculate_renewal_price(subscription):
                        results["insufficient_balance"] += 1
                    else:
                        results["failed"] += 1
            except Exception:
                logger.exception("Error processing autorenewal for user %d", subscription.user_id)
                results["failed"] += 1

        logger.info(
            "Autorenewal completed: success=%d, failed=%d, insufficient_balance=%d",
            results["success"],
            results["failed"],
            results["insufficient_balance"],
        )

        return results