malenia-backend/services/subscriptions.py

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 external.rw import sync_subscription_by_telegram_id
from schemas.enums import SubscriptionStatus
from schemas.plans import PricingPlans


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
) -> 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,
    )
    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)
    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 sync_subscription_by_telegram_id(
            expires_at=subscription.expires_at,
            devices=subscription.devices,
            telegram_id=user.telegram_id,
            username=user.username,
        )