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.types import LabeledPrice, Message, PreCheckoutQuery, SuccessfulPayment

from user_store import (
    SUBSCRIPTION_PRICE_STARS,
    activate_subscription,
    connect,
    ensure_subscription,
)

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

STARS_CURRENCY = "XTR"
INVOICE_PREFIX = "sub30:"


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) -> None:
    """Отправляет инвойс на 30 дней подписки за Stars."""
    payload = f"{INVOICE_PREFIX}{datetime.datetime.now(datetime.timezone.utc):%Y%m%d%H%M%S}"
    kwargs: dict[str, object] = {
        "chat_id": chat_id,
        "title": "Подписка на 30 дней",
        "description": "Доступ к генерации постов, идей и голосовым на 30 дней.",
        "payload": payload,
        "currency": STARS_CURRENCY,
        "prices": [LabeledPrice(label="30 дней", amount=SUBSCRIPTION_PRICE_STARS)],
        "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:
    """Подтверждаем только свой инвойс, в правильной валюте и по правильной цене."""
    if query.currency != STARS_CURRENCY or query.total_amount != SUBSCRIPTION_PRICE_STARS:
        await query.answer(ok=False, error_message="Некорректная сумма или валюта платежа.")
        return
    if not query.invoice_payload.startswith(INVOICE_PREFIX):
        await query.answer(ok=False, error_message="Этот счёт не принадлежит подписке.")
        return
    await query.answer(ok=True)


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

    charge_id = payment.telegram_payment_charge_id
    if payment.currency != STARS_CURRENCY or payment.total_amount != SUBSCRIPTION_PRICE_STARS:
        await _reject(message, charge_id, "неверная сумма или валюта")
        return
    if not payment.invoice_payload.startswith(INVOICE_PREFIX):
        await _reject(message, charge_id, "неизвестный invoice_payload")
        return

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

    await asyncio.to_thread(ensure_subscription, message.from_user.id)
    days = await asyncio.to_thread(activate_subscription, message.from_user.id, months=1)
    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(
        f"Оплата получена — подписка активирована на {days} дн.\n\n"
        f"Списано {payment.total_amount} звёзд. Теперь доступны генерация постов, "
        f"идей и голосовые. Команда /balance покажет статус."
    )


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(
        "Платёж получен, но активировать подписку не получилось. Напиши, пожалуйста, поддержке."
    )