Files
Solo_bot/handlers/refferal.py
T

305 lines
12 KiB
Python

import os
from io import BytesIO
import qrcode
from aiogram import F, Router
from aiogram.enums import ParseMode
from aiogram.fsm.context import FSMContext
from aiogram.types import (
CallbackQuery,
InlineKeyboardButton,
InlineQuery,
InlineQueryResultArticle,
InputTextMessageContent,
Message,
)
from aiogram.utils.keyboard import InlineKeyboardBuilder
from sqlalchemy import desc, func, select
from sqlalchemy.ext.asyncio import AsyncSession
from bot import bot
from config import ADMIN_ID, INLINE_MODE, REFERRAL_BONUS_PERCENTAGES, REFERRAL_QR, TOP_REFERRAL_BUTTON, USERNAME_BOT
from core.bootstrap import BUTTONS_CONFIG, MODES_CONFIG
from database import (
add_referral,
add_user,
get_referral_by_referred_id,
get_referral_stats,
)
from database.models import Referral
from database.tariffs import get_tariffs
from handlers.buttons import BACK, INVITE, MAIN_MENU, QR, TOP_FIVE
from handlers.payments.currency_rates import format_for_user
from handlers.texts import (
INVITE_MESSAGE_TEMPLATE,
INVITE_TEXT_NON_INLINE,
NEW_REFERRAL_NOTIFICATION,
REFERRAL_OFFERS,
REFERRAL_SUCCESS_MSG,
TOP_REFERRALS_TEXT,
)
from logger import logger
from .texts import get_referral_link
from .utils import edit_or_send_message, format_days
router = Router()
@router.callback_query(F.data == "invite")
@router.message(F.text == "/invite")
async def invite_handler(callback_query_or_message: Message | CallbackQuery, session: AsyncSession):
if isinstance(callback_query_or_message, CallbackQuery):
chat_id = callback_query_or_message.message.chat.id
target_message = callback_query_or_message.message
language_code = callback_query_or_message.from_user.language_code
else:
chat_id = callback_query_or_message.chat.id
target_message = callback_query_or_message
language_code = callback_query_or_message.from_user.language_code
referral_link = get_referral_link(chat_id)
referral_stats = await get_referral_stats(session, chat_id)
bonuses_lines = []
for level, value in REFERRAL_BONUS_PERCENTAGES.items():
if isinstance(value, float):
bonuses_lines.append(f"{level} уровень: 🌟 {int(value * 100)}% бонуса")
else:
value_text = await format_for_user(session, chat_id, value, language_code)
bonuses_lines.append(f"{level} уровень: 💸 {value_text} бонуса")
bonuses_block = "\n".join(bonuses_lines)
details_lines = []
for level, stats in referral_stats["referrals_by_level"].items():
bonus_value = REFERRAL_BONUS_PERCENTAGES.get(level)
if isinstance(bonus_value, float):
bonus_text = f"{int(bonus_value * 100)}%"
else:
bonus_text = await format_for_user(session, chat_id, bonus_value, language_code)
details_lines.append(f"🔹 Уровень {level}: {stats['total']} - {bonus_text}")
details_block = "\n".join(details_lines)
total_bonus_text = await format_for_user(session, chat_id, referral_stats["total_referral_bonus"], language_code)
invite_message = INVITE_MESSAGE_TEMPLATE.format(
referral_link=referral_link,
bonuses_block=bonuses_block,
total_referrals=referral_stats["total_referrals"],
details_block=details_block,
total_referral_bonus=total_bonus_text,
)
image_path = os.path.join("img", "pic_invite.jpg")
inline_mode_enabled = bool(MODES_CONFIG.get("INLINE_MODE_ENABLED", INLINE_MODE))
builder = InlineKeyboardBuilder()
if inline_mode_enabled:
builder.button(text=INVITE, switch_inline_query="invite")
else:
invite_text = INVITE_TEXT_NON_INLINE.format(referral_link=referral_link)
builder.button(text=INVITE, switch_inline_query=invite_text)
if BUTTONS_CONFIG.get("REFERRAL_QR_BUTTON_ENABLE", REFERRAL_QR):
builder.button(text=QR, callback_data=f"show_referral_qr|{chat_id}")
if BUTTONS_CONFIG.get("TOP_REFERRAL_BUTTON_ENABLE", TOP_REFERRAL_BUTTON):
builder.button(text=TOP_FIVE, callback_data="top_referrals")
builder.button(text=MAIN_MENU, callback_data="profile")
builder.adjust(1)
await edit_or_send_message(
target_message=target_message,
text=invite_message,
reply_markup=builder.as_markup(),
media_path=image_path,
disable_web_page_preview=False,
)
@router.inline_query(F.query.in_(["referral", "ref", "invite"]))
async def inline_referral_handler(inline_query: InlineQuery, session: AsyncSession):
referral_link = f"https://t.me/{USERNAME_BOT}?start=referral_{inline_query.from_user.id}"
trial_tariffs = await get_tariffs(session, group_code="trial")
if not trial_tariffs:
await inline_query.answer(results=[], cache_time=0)
return
trial_days = trial_tariffs[0]["duration_days"]
trial_time_formatted = format_days(trial_days)
results: list[InlineQueryResultArticle] = []
for index, offer in enumerate(REFERRAL_OFFERS):
message_text = offer["message"].format(trial_time_formatted=trial_time_formatted)[:4096]
title = offer["title"].format(trial_time_formatted=trial_time_formatted)
description = offer["description"]
builder = InlineKeyboardBuilder()
builder.row(InlineKeyboardButton(text=title, url=referral_link))
results.append(
InlineQueryResultArticle(
id=str(index),
title=title,
description=description,
input_message_content=InputTextMessageContent(
message_text=message_text,
parse_mode=ParseMode.HTML,
),
reply_markup=builder.as_markup(),
)
)
await inline_query.answer(results=results, cache_time=60, is_personal=True)
@router.callback_query(F.data.startswith("show_referral_qr|"))
async def show_referral_qr(callback_query: CallbackQuery):
try:
chat_id = callback_query.data.split("|")[1]
referral_link = get_referral_link(chat_id)
qr = qrcode.QRCode(version=1, box_size=10, border=4)
qr.add_data(referral_link)
qr.make(fit=True)
image = qr.make_image(fill_color="black", back_color="white")
buffer = BytesIO()
image.save(buffer, format="PNG")
buffer.seek(0)
qr_path = f"/tmp/qrcode_referral_{chat_id}.png"
with open(qr_path, "wb") as file:
file.write(buffer.read())
builder = InlineKeyboardBuilder()
builder.row(InlineKeyboardButton(text=BACK, callback_data="invite"))
builder.row(InlineKeyboardButton(text=MAIN_MENU, callback_data="profile"))
await edit_or_send_message(
target_message=callback_query.message,
text="📷 <b>Ваш QR-код для реферальной ссылки.</b>",
reply_markup=builder.as_markup(),
media_path=qr_path,
)
os.remove(qr_path)
except Exception as error:
logger.error(f"Ошибка при генерации QR-кода для реферальной ссылки: {error}", exc_info=True)
await callback_query.message.answer("❌ Произошла ошибка при создании QR-кода.")
@router.callback_query(F.data == "top_referrals")
async def top_referrals_handler(callback_query: CallbackQuery, session: AsyncSession):
user_id = callback_query.from_user.id
result = await session.execute(select(func.count()).select_from(Referral).where(Referral.referrer_tg_id == user_id))
user_referral_count = result.scalar_one() or 0
personal_block = "Твоё место в рейтинге:\n"
if user_referral_count > 0:
subquery = (
select(func.count().label("cnt"))
.select_from(Referral)
.group_by(Referral.referrer_tg_id)
.having(func.count() > user_referral_count)
.subquery()
)
result = await session.execute(select(func.count()).select_from(subquery))
user_position = result.scalar_one() + 1
personal_block += f"{user_position}. {user_id} - {user_referral_count} чел."
else:
personal_block += "Ты еще не приглашал пользователей в проект."
result = await session.execute(
select(
Referral.referrer_tg_id,
func.count(Referral.referred_tg_id).label("referral_count"),
)
.group_by(Referral.referrer_tg_id)
.order_by(desc("referral_count"))
.limit(5)
)
top_referrals = result.all()
is_admin = user_id in ADMIN_ID
rows = ""
for index, row in enumerate(top_referrals, 1):
referrer_id = str(row.referrer_tg_id)
count = row.referral_count
display_id = referrer_id if is_admin else f"{referrer_id[:5]}*****"
rows += f"{index}. {display_id} - {count} чел.\n"
text = TOP_REFERRALS_TEXT.format(personal_block=personal_block, rows=rows)
builder = InlineKeyboardBuilder()
builder.row(InlineKeyboardButton(text=BACK, callback_data="invite"))
builder.row(InlineKeyboardButton(text=MAIN_MENU, callback_data="profile"))
await edit_or_send_message(
target_message=callback_query.message,
text=text,
reply_markup=builder.as_markup(),
media_path=None,
disable_web_page_preview=False,
)
async def handle_referral_link(
referral_code: str,
message: Message,
state: FSMContext,
session: AsyncSession,
user_data: dict | None = None,
):
try:
referrer_tg_id = int(referral_code)
user = user_data or message.from_user or message.chat
user_id = user["tg_id"] if isinstance(user, dict) else user.id
if referrer_tg_id == user_id:
await message.answer("❌ Вы не можете быть реферальной ссылкой самого себя.")
return
existing_referral = await get_referral_by_referred_id(session, user_id)
if existing_referral:
await message.answer("❌ Вы уже использовали реферальную ссылку.")
return
if isinstance(user, dict):
inserted = await add_user(session=session, **user)
else:
inserted = await add_user(
session=session,
tg_id=user.id,
username=getattr(user, "username", None),
first_name=getattr(user, "first_name", None),
last_name=getattr(user, "last_name", None),
language_code=getattr(user, "language_code", None),
is_bot=getattr(user, "is_bot", False),
)
if not inserted:
await message.answer("❌ Вы уже зарегистрированы и не можете стать рефералом.")
return
await add_referral(session, user_id, referrer_tg_id)
try:
await bot.send_message(
referrer_tg_id,
NEW_REFERRAL_NOTIFICATION.format(referred_id=user_id),
)
except Exception as error:
logger.error(f"Не удалось отправить уведомление пригласившему ({referrer_tg_id}): {error}")
await message.answer(REFERRAL_SUCCESS_MSG.format(referrer_tg_id=referrer_tg_id))
except Exception as error:
logger.error(f"Ошибка при обработке реферальной ссылки {referral_code}: {error}")
await message.answer("❌ Произошла ошибка при обработке реферальной ссылки.")