malenia/bot/handlers/buy.py

import logging
import math

from aiogram import F, Router
from aiogram.fsm.context import FSMContext
from aiogram.types import CallbackQuery

from bot.keyboards.client import (
    devices_selector,
    duration_selector,
    payment_gateways,
    return_to_menu,
)
from bot.states.buy import BillingStorage, SubscriptionStorage
from bot.texts import (
    CALLBACK_FALLBACK,
    CHECKOUT_PAYMENT_CHOICE,
    CONTEXT_REINITIATED,
    NOTHING_PLACEHOLDER,
    STANDARD_FALLBACK,
    SUBSCRIPTION_DEVICE_SELECTOR,
    SUBSCRIPTION_DURATION_SELECTOR,
)
from config import settings
from misc.utils import (
    calculate_price,
    convert_duration_to_int,
    convert_int_to_duration,
    get_discount,
)
from schemas.billing import BillingContext
from schemas.subscriptions import SubscriptionPlan

router = Router()
logger = logging.getLogger(__name__)


@router.callback_query(F.data == "buy")
async def buy_init(cb: CallbackQuery, state: FSMContext):
    await state.clear()

    ctx = SubscriptionPlan()
    await cb.message.edit_text(
        SUBSCRIPTION_DEVICE_SELECTOR.format(price=calculate_price(ctx)),
        reply_markup=devices_selector(settings.min_devices),
    )

    await state.set_state(SubscriptionStorage.devices)
    await state.set_data({"ctx": ctx})


@router.callback_query(F.data.startswith("devices:"))
async def device_count_select(cb: CallbackQuery, state: FSMContext):
    data = await state.get_data()
    await state.clear()

    ctx: SubscriptionPlan | None = data.get("ctx", SubscriptionPlan())

    cb_data = cb.data.split(":")
    devices = int(cb_data[1]) if cb_data[1].isdigit() else settings.min_devices
    ctx.devices = devices
    ctx.whitelists = ctx.whitelists or devices >= settings.whitelist_device_threshold

    await cb.message.edit_text(
        SUBSCRIPTION_DEVICE_SELECTOR.format(price=calculate_price(ctx)),
        reply_markup=devices_selector(current=ctx.devices, whitelists=ctx.whitelists),
    )
    await state.set_state(SubscriptionStorage.devices)
    await state.set_data({"ctx": ctx})


@router.callback_query(SubscriptionStorage.devices, F.data == "whitelist:toggle")
async def toggle_whitelists(cb: CallbackQuery, state: FSMContext):
    data = await state.get_data()
    await state.clear()

    ctx: SubscriptionPlan | None = data.get("ctx", SubscriptionPlan())
    ctx.whitelists = not ctx.whitelists or ctx.devices >= settings.whitelist_device_threshold

    try:
        await cb.message.edit_text(
            SUBSCRIPTION_DEVICE_SELECTOR.format(price=calculate_price(ctx)),
            reply_markup=devices_selector(current=ctx.devices, whitelists=ctx.whitelists),
        )
    except Exception as exc:
        if "message is not modified:" in str(exc):
            await cb.answer(NOTHING_PLACEHOLDER)
        else:
            await cb.answer(CALLBACK_FALLBACK)
    await state.set_state(SubscriptionStorage.devices)
    await state.set_data({"ctx": ctx})


@router.callback_query(SubscriptionStorage.devices, F.data == "confirm")
async def select_duration_init(cb: CallbackQuery, state: FSMContext):
    data = await state.get_data()
    await state.clear()

    ctx: SubscriptionPlan | None = data.get("ctx")
    if not ctx:
        await cb.answer(CALLBACK_FALLBACK)
        return

    try:
        discount = get_discount(convert_duration_to_int(ctx.duration))
        ctx.discount = 1 - discount / 100
        await cb.message.edit_text(
            SUBSCRIPTION_DURATION_SELECTOR.format(
                price=calculate_price(ctx),
                discount=math.floor(discount),
            ),
            reply_markup=duration_selector(current=ctx.duration, back_cb=f"devices:{ctx.devices}"),
        )
    except Exception as exc:
        if "message is not modified:" in str(exc):
            await cb.answer(NOTHING_PLACEHOLDER)
        else:
            await cb.answer(CALLBACK_FALLBACK)
            raise

    await state.set_state(SubscriptionStorage.duration)
    await state.set_data({"ctx": ctx})


@router.callback_query(F.data.startswith("duration:"))
async def select_duration(cb: CallbackQuery, state: FSMContext):
    data = await state.get_data()
    await state.clear()

    ctx: SubscriptionPlan | None = data.get("ctx")
    if not ctx:
        await cb.answer(CONTEXT_REINITIATED)
        ctx = SubscriptionPlan()

    cb_data = cb.data.split(":")
    duration = int(cb_data[1]) if cb_data[1].isdigit() else 1
    ctx.duration = convert_int_to_duration(duration)
    try:
        discount = get_discount(convert_duration_to_int(ctx.duration))
        ctx.discount = 1 - discount / 100
        await cb.message.edit_text(
            SUBSCRIPTION_DURATION_SELECTOR.format(
                price=calculate_price(ctx),
                discount=math.floor(discount),
            ),
            reply_markup=duration_selector(current=ctx.duration, back_cb=f"devices:{ctx.devices}"),
        )
    except Exception as exc:
        if "message is not modified:" in str(exc):
            await cb.answer(NOTHING_PLACEHOLDER)
        else:
            await cb.answer(CALLBACK_FALLBACK)
            raise

    await state.set_state(SubscriptionStorage.duration)
    await state.set_data({"ctx": ctx})


@router.callback_query(SubscriptionStorage.duration, F.data == "confirm")
async def proceed_to_checkout(cb: CallbackQuery, state: FSMContext):
    data = await state.get_data()
    await state.clear()

    ctx: SubscriptionPlan | None = data.get("ctx")
    if not ctx:
        await cb.message.edit_text(STANDARD_FALLBACK, reply_markup=return_to_menu)
        return

    await cb.message.edit_text(
        CHECKOUT_PAYMENT_CHOICE,
        reply_markup=payment_gateways,
    )  # FIXME: temp solution that's incredibly stupid

    await state.set_state(BillingStorage.pending)
    billing_context = BillingContext(amount=calculate_price(ctx))
    await state.set_data({"ctx": billing_context})