social-post-writer-bot/payments.py

"""Реальная оплата Telegram Stars (XTR) за подписку.

Цикл: sendInvoice -> pre_checkout_query -> answerPreCheckoutQuery ->
successful_payment -> activate_subscription.
`telegram_payment_charge_id` каждого платежа пишется в таблицу payments,
поэтому один платёж не может активировать подписку дважды, а по нему потом
можно сделать возврат через refundStarPayment.
"""
from __future__ import annotations

import asyncio
import datetime
import logging
import os
from contextlib import closing

from aiogram import F, Router
from aiogram.fsm.context import FSMContext
from aiogram.types import LabeledPrice, Message, PreCheckoutQuery, SuccessfulPayment

from keyboards import main_menu_keyboard

from user_store import (
    SUBSCRIPTION_PLANS,
    activate_subscription_days,
    activate_subscription,
    connect,
    ensure_subscription,
    subscription_row,
    subscription_status_text,
)

logger = logging.getLogger(__name__)
router = Router(name="payments")

STARS_CURRENCY = "XTR"
INVOICE_PREFIX = "sub:"


def plan_for_payload(payload: str) -> tuple[int, int] | None:
    if payload.startswith("sub30:"):
        try:
            timestamp = payload.split(":", 1)[1]
            datetime.datetime.strptime(timestamp, "%Y%m%d%H%M%S")
        except (ValueError, IndexError):
            return None
        return (30, SUBSCRIPTION_PLANS[30])
    try:
        prefix, days_part, timestamp = payload.split(":", 2)
        if prefix != "sub" or not timestamp.isdigit():
            return None
        days = int(days_part)
    except (ValueError, IndexError):
        return None
    price = SUBSCRIPTION_PLANS.get(days)
    if price is None or (days == 0 and price != 1):
        return None
    return days, price


def _db_path(db_path: str | None = None) -> str | None:
    return db_path if db_path is not None else os.getenv("BOT_DB_PATH")


def init_payments_table(db_path: str | None = None) -> None:
    with closing(connect(db_path)) as connection, connection:
        connection.execute(
            """
            CREATE TABLE IF NOT EXISTS payments (
                telegram_payment_charge_id TEXT PRIMARY KEY,
                telegram_user_id INTEGER NOT NULL,
                invoice_payload TEXT NOT NULL,
                currency TEXT NOT NULL,
                total_amount INTEGER NOT NULL,
                created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
            )
            """
        )
        connection.execute(
            "CREATE INDEX IF NOT EXISTS idx_payments_user ON payments(telegram_user_id)"
        )


def payment_recorded(charge_id: str, db_path: str | None = None) -> bool:
    with closing(connect(db_path)) as connection:
        row = connection.execute(
            "SELECT 1 FROM payments WHERE telegram_payment_charge_id = ?", (charge_id,)
        ).fetchone()
    return row is not None


def save_payment(
    charge_id: str,
    user_id: int,
    invoice_payload: str,
    currency: str,
    total_amount: int,
    db_path: str | None = None,
) -> None:
    with closing(connect(db_path)) as connection, connection:
        connection.execute(
            """
            INSERT OR IGNORE INTO payments
                (telegram_payment_charge_id, telegram_user_id, invoice_payload, currency, total_amount)
            VALUES (?, ?, ?, ?, ?)
            """,
            (charge_id, user_id, invoice_payload, currency, total_amount),
        )


async def send_subscription_invoice(
    bot, chat_id: int, thread_id: int | None = None, days: int = 30
) -> None:
    """Отправляет выбранный инвойс на разовую подписку в Stars."""
    price = SUBSCRIPTION_PLANS[days]
    payload = f"{INVOICE_PREFIX}{days}:{datetime.datetime.now(datetime.timezone.utc):%Y%m%d%H%M%S}"
    kwargs: dict[str, object] = {
        "chat_id": chat_id,
        "title": f"Подписка на {days} дней",
        "description": "Доступ к боту по генерации текста постов",
        "payload": payload,
        "currency": STARS_CURRENCY,
        "prices": [LabeledPrice(label=f"{days} дней", amount=price)],
        "provider_token": "",
    }
    if thread_id is not None:
        kwargs["message_thread_id"] = thread_id
    await bot.send_invoice(**kwargs)


@router.pre_checkout_query()
async def handle_pre_checkout(query: PreCheckoutQuery) -> None:
    """Подтверждаем только свой инвойс, в правильной валюте и по правильной цене."""
    plan = plan_for_payload(query.invoice_payload)
    if plan is None or query.currency != STARS_CURRENCY or query.total_amount != plan[1]:
        await query.answer(ok=False, error_message="Некорректная сумма или валюта платежа.")
        return
    await query.answer(ok=True)


async def _return_to_main_menu(message: Message, state: FSMContext | None) -> None:
    """После успешной оплаты/продления возвращает пользователя в главное меню."""
    if state is not None:
        await state.clear()
        await state.set_state("BotMode:lobby")
    await message.answer("Оплата принята.", reply_markup=main_menu_keyboard())


@router.message(F.successful_payment)
async def handle_successful_payment(message: Message, state: FSMContext | None = None) -> None:
    payment = message.successful_payment
    if payment is None or message.from_user is None:
        return

    charge_id = payment.telegram_payment_charge_id
    plan = plan_for_payload(payment.invoice_payload)
    if plan is None or payment.currency != STARS_CURRENCY or payment.total_amount != plan[1]:
        await _reject(message, charge_id, "неверная сумма или валюта")
        return
    days, _price = plan

    db_path = _db_path()
    if await asyncio.to_thread(payment_recorded, charge_id, db_path):
        await message.answer("Этот платёж уже учтён — подписка не продлена повторно.")
        return

    if days == 0:
        days = 1

    await asyncio.to_thread(ensure_subscription, message.from_user.id)
    await asyncio.to_thread(activate_subscription_days, message.from_user.id, days)
    sub = await asyncio.to_thread(subscription_row, message.from_user.id, db_path)
    await asyncio.to_thread(
        save_payment,
        charge_id,
        message.from_user.id,
        payment.invoice_payload,
        payment.currency,
        payment.total_amount,
        db_path,
    )
    await message.answer(subscription_status_text(sub))
    await _return_to_main_menu(message, state)


async def _reject(message: Message, charge_id: str, reason: str) -> None:
    logger.error(
        "Rejected payment charge=%s user=%s reason=%s",
        charge_id,
        message.from_user.id if message.from_user else None,
        reason,
    )
    await message.answer(
        "Платёж получен, но активировать подписку не получилось. Напиши, пожалуйста, поддержке."
    )