from datetime import UTC, datetime, timedelta from sqlalchemy import delete from sqlalchemy.ext.asyncio import AsyncSession from db.models import Subscription, SubscriptionAddon, User from db.models.orders import Order, OrderStatus from db.models.transactions import BalanceTransaction, BalanceTxType from schemas.enums import SubscriptionStatus from schemas.plans import PricingPlans from services.rw_sync import enqueue_rw_sync def calculate_plan_monthly_price( pricing: PricingPlans, devices: int, addon_ids: list[str] ) -> float: addon_prices = {addon.id: addon.price for addon in pricing.addons} return devices * pricing.device_price + sum(addon_prices[addon_id] for addon_id in addon_ids) def calculate_order_total( pricing: PricingPlans, devices: int, addon_ids: list[str], duration_days: int ) -> float: monthly_price = calculate_plan_monthly_price(pricing, devices, addon_ids) return monthly_price * duration_days / 30 def get_subscription_addon_ids(subscription: Subscription | None) -> list[str]: if subscription is None: return [] return [addon.addon_id for addon in subscription.addons] def should_apply_immediately( *, order: Order, pricing: PricingPlans, now: datetime, subscription: Subscription | None = None, ) -> bool: if subscription is None or subscription.expires_at <= now: return True current_addons = set(get_subscription_addon_ids(subscription)) new_addons = {addon.addon_id for addon in order.addons} if order.devices == subscription.devices and new_addons == current_addons: return True current_monthly_price = calculate_plan_monthly_price( pricing, subscription.devices, list(current_addons) ) new_monthly_price = calculate_plan_monthly_price(pricing, order.devices, list(new_addons)) return ( order.devices >= subscription.devices and new_addons.issuperset(current_addons) and new_monthly_price >= current_monthly_price ) async def replace_subscription_addons( session: AsyncSession, subscription_id: int, addon_ids: list[str] ) -> None: await session.execute( delete(SubscriptionAddon).where(SubscriptionAddon.subscription_id == subscription_id) ) for addon_id in addon_ids: session.add(SubscriptionAddon(subscription_id=subscription_id, addon_id=addon_id)) async def ensure_subscription( user: User, session: AsyncSession, starts_at: datetime, duration_days: int ) -> Subscription: subscription = user.subscription if subscription is not None: return subscription subscription = Subscription( user_id=user.id, devices=0, status=SubscriptionStatus.EXPIRED, expires_at=starts_at, duration_days=duration_days, ) session.add(subscription) await session.flush() user.subscription = subscription return subscription async def deduct_order_balance( session: AsyncSession, *, user: User, order: Order, description: str, ) -> None: if order.balance_amount <= 0: return session.add( BalanceTransaction( user_id=user.id, amount=-order.balance_amount, tx_type=BalanceTxType.PURCHASE, balance_before=user.balance, balance_after=user.balance - order.balance_amount, description=description, ) ) user.balance -= order.balance_amount async def apply_order_now( session: AsyncSession, *, user: User, order: Order, pricing: PricingPlans, now: datetime, ) -> Subscription: subscription = await ensure_subscription(user, session, now, order.duration_days) addon_ids = [addon.addon_id for addon in order.addons] current_addons = [] if subscription.devices == 0 else get_subscription_addon_ids(subscription) if subscription.status == SubscriptionStatus.ACTIVE and subscription.expires_at > now: if order.devices == subscription.devices and set(addon_ids) == set(current_addons): subscription.expires_at += timedelta(days=order.duration_days) else: current_monthly_price = calculate_plan_monthly_price( pricing, subscription.devices, current_addons ) next_monthly_price = calculate_plan_monthly_price(pricing, order.devices, addon_ids) remaining_seconds = (subscription.expires_at - now).total_seconds() remaining_days = max(remaining_seconds / 86400, 0) remaining_credit = current_monthly_price * remaining_days / 30 purchased_value = calculate_order_total( pricing, order.devices, addon_ids, order.duration_days ) total_days = ((remaining_credit + purchased_value) / next_monthly_price) * 30 subscription.expires_at = now + timedelta(days=total_days) else: subscription.expires_at = now + timedelta(days=order.duration_days) subscription.devices = order.devices subscription.status = SubscriptionStatus.ACTIVE await replace_subscription_addons(session, subscription.id, addon_ids) order.applies_at = now order.applied_at = now return subscription async def queue_order_for_later(*, order: Order, subscription: Subscription, now: datetime) -> None: order.applies_at = max(subscription.expires_at, now) async def sync_user_subscription( session: AsyncSession, *, user: User, ) -> None: now = datetime.now(UTC) subscription = user.subscription applied_due_orders = False if subscription is not None and subscription.expires_at <= now: due_orders = [ order for order in user.orders if order.status == OrderStatus.PAID and order.applied_at is None and order.applies_at is not None and order.applies_at <= now ] due_orders.sort(key=lambda order: (order.applies_at, order.id)) if due_orders: for order in due_orders: base_time = max(subscription.expires_at, order.applies_at) subscription.devices = order.devices subscription.status = SubscriptionStatus.ACTIVE subscription.expires_at = base_time + timedelta(days=order.duration_days) await replace_subscription_addons( session, subscription.id, [addon.addon_id for addon in order.addons], ) order.applied_at = now applied_due_orders = True else: subscription.status = SubscriptionStatus.EXPIRED if applied_due_orders: await enqueue_rw_sync(session, subscription.id)