Referral-system

pyproject.toml

@@ -10,4 +10,5 @@ dependencies = [
     "dotenv>=0.9.9",
     "pydantic>=2.11.10",
     "pymongo>=4.15.4",
+    "qrcode[pil]>=8.2",
 ]

requirements.txt

@@ -2,4 +2,6 @@ aiogram>=3.22.0
 asyncio>=4.0.0
 dotenv>=0.9.9
 pydantic>=2.11.10
-pymongo>=4.15.4
+pymongo>=4.15.4
+pymongo>=4.15.4
+qrcode[pil]>=8.2

src/bot.py

@@ -1,22 +1,30 @@
+import io
 import os
+from typing import List
 
+from aiogram.utils.payload import decode_payload
+import qrcode
 from aiogram import Bot, Dispatcher, F, types
 from aiogram.client.default import DefaultBotProperties
 from aiogram.enums import ParseMode
-from aiogram.filters import CommandStart
+from aiogram.filters import CommandObject, CommandStart
 from aiogram.fsm.context import FSMContext
 from aiogram.fsm.state import State, StatesGroup
 from aiogram.fsm.storage.pymongo import PyMongoStorage
 from aiogram.types import (
+    BufferedInputFile,
     CallbackQuery,
     InlineKeyboardButton,
     InlineKeyboardMarkup,
     KeyboardButton,
+    ReactionTypeEmoji,
     ReplyKeyboardMarkup,
 )
+from aiogram.utils.deep_linking import create_start_link
 from aiogram.utils.formatting import as_list
 from dotenv import load_dotenv
 from pymongo import AsyncMongoClient
+from qrcode import QRCode
 
 from . import callbacks, templates
 from .models import Link, User, check_tg_username
@@ -43,21 +51,38 @@ async def main():
 
 rkb = ReplyKeyboardMarkup(
     keyboard=[
-        [KeyboardButton(text="Мои контакты")],
-        [KeyboardButton(text="Узнать тип личности")],
-        [KeyboardButton(text="Кол-во пользователей")],
+        [KeyboardButton(text="Мои контакты"), KeyboardButton(text="Тип личности")],
+        [
+            KeyboardButton(text="Кол-во пользователей"),
+            KeyboardButton(text="Реферальная система"),
+        ],
     ]
 )
 
 
 class AddingUser(StatesGroup):
+    starting = State()
     sex = State()
     course = State()
     living_place = State()
 
 
 @dp.message(CommandStart())
-async def start_handler(message: types.Message, state: FSMContext):
+async def start_handler(
+    message: types.Message, command: CommandObject, state: FSMContext
+):
+    await state.set_state(AddingUser.starting)
+    if command.args:
+        linked_by = decode_payload(command.args)
+        if linked_by != message.from_user.username:
+            user = await userdb.get_user(message.from_user.username)
+            if user is None:
+                await state.update_data(invited_by=linked_by)
+                await userdb.add_invited(linked_by, message.from_user.username)
+            elif len(user.links) < 5 and user.invited_by is None:
+                await userdb.add_invited_by(message.from_user.username, linked_by)
+                await userdb.add_invited(linked_by, message.from_user.username)
+
     await message.answer(
         templates.starting_message,
         reply_markup=InlineKeyboardMarkup(
@@ -189,9 +214,12 @@ async def process_living(
     await userdb.add_user(
         User(
             username=query.from_user.username,
+            userid=query.from_user.id,
+            chatid=query.message.chat.id,
             sex=state_data["sex"],
             course=state_data["course"],
             living=callback_data.living,
+            invited_by=state_data.get("invited_by"),
         )
     )
     await state.clear()
@@ -231,6 +259,76 @@ def rating_to_text(rating: int) -> str:
     raise Exception("Rating_to_text получил invalid значение")
 
 
+@dp.message(F.text[0] == "@")
+async def user_name_checker(message: types.Message):
+    await userdb.add_ids_to_user(
+        message.from_user.username, message.from_user.id, message.chat.id
+    )
+    msg = (message.text).strip()
+    try:
+        username_to = check_tg_username(msg)
+    except ValueError:
+        await message.answer('Напиши юзернейм в формате "@username"')
+        return
+    if username_to == message.from_user.username:
+        await message.react([ReactionTypeEmoji(emoji="🥰")])
+        await message.answer(
+            "Любовь к себе это, конечно, хорошо, но, пожалуйста, добавь кого-нибудь другого"
+        )
+        return
+
+    kb = InlineKeyboardMarkup(
+        inline_keyboard=[
+            [
+                InlineKeyboardButton(
+                    text="Близкий друг",
+                    callback_data=callbacks.LinkCallback(
+                        username_to=username_to, rating=3
+                    ).pack(),
+                )
+            ],
+            [
+                InlineKeyboardButton(
+                    text="Приятель",
+                    callback_data=callbacks.LinkCallback(
+                        username_to=username_to, rating=2
+                    ).pack(),
+                ),
+            ],
+            [
+                InlineKeyboardButton(
+                    text="Знакомый",
+                    callback_data=callbacks.LinkCallback(
+                        username_to=username_to, rating=1
+                    ).pack(),
+                ),
+            ],
+        ]
+    )
+
+    await message.answer("Кто он для тебя?", reply_markup=kb)
+
+
+@dp.callback_query(callbacks.LinkCallback.filter())
+async def process_data(query: CallbackQuery, callback_data: callbacks.LinkCallback):
+    from_username = query.from_user.username
+    if await userdb.get_user(from_username) is None:
+        await query.message.answer("Похоже тебе нужно перезапустить бота: /start")
+        return
+    await userdb.add_link(
+        from_username,
+        Link(username_to=callback_data.username_to, rating=callback_data.rating),
+    )
+    await query.message.edit_text(
+        as_list(
+            f"✅ @{callback_data.username_to} добавлен как {rating_to_text(callback_data.rating).lower()}",
+            "\n📝 Чтобы добавить ещё друга — просто введи следующий юзернейм.",
+            "\n🔁 Чем больше друзей ты добавишь — тем точнее будет твой социальный портрет!",
+        ).as_html(),
+        reply_markup=None,
+    )
+
+
 @dp.message(F.text == "Мои контакты")
 async def get_usS(message: types.Message):
     await userdb.add_ids_to_user(
@@ -253,21 +351,21 @@ async def get_usS(message: types.Message):
     await message.answer(all_users_and_rating, reply_markup=rkb)
 
 
-@dp.message(F.text == "Узнать тип личности")
+@dp.message(F.text == "Тип личности")
 async def get_summary(message: types.Message):
     await userdb.add_ids_to_user(
         message.from_user.username, message.from_user.id, message.chat.id
     )
     links = await userdb.get_links(message.from_user.username)
     ratings = [i.rating for i in links]
-    p1 = ratings.count(1) / len(ratings)
-    p2 = ratings.count(2) / len(ratings)
-    p3 = ratings.count(3) / len(ratings)
     if len(ratings) < 5:
         await message.answer(
             "К сожалению, ты написал слишком мало для полноценного отчёта. Давай постараемся добавить всех друзей!"
         )
-    elif p3 >= 0.5:
+    p1 = ratings.count(1) / len(ratings)
+    p2 = ratings.count(2) / len(ratings)
+    p3 = ratings.count(3) / len(ratings)
+    if p3 >= 0.5:
         await message.answer(
             templates.make_type_str(
                 "Сердце компании",
@@ -354,65 +452,64 @@ async def get_count(message: types.Message):
     )
 
 
-@dp.message(F.text[0] == "@")
-async def user_name_checker(message: types.Message):
-    userdb.add_ids_to_user(
+@dp.message(F.text == "Реферальная система")
+async def get_referral(message: types.Message):
+    def generate_message(
+        link: str, str_list: List[str] = None, points: int = None
+    ) -> str:
+        message = (
+            "**🚀 Участвуй в турнире с реферальной системой\\!**\n"
+            + f"Твоя личная ссылка для приглашений:\n`{link}`\n\\(Нажми чтобы скопировать\\)\n\n"
+            + "✨ Каждый приглашённый друг \\= \\+1 к твоим шансам на победу\\!\n"
+            + "Главное — чтобы он указал минимум 5 связей\n"
+            + "Приглашение засчитывается, если у человека заполнено меньше 5 связей"
+        )
+
+        if not message:
+            message += "Пока что никто не переходил по ссылке"
+        return message
+
+    await userdb.add_ids_to_user(
         message.from_user.username, message.from_user.id, message.chat.id
     )
-    msg = (message.text).strip()
-    try:
-        username_to = check_tg_username(msg)
-    except ValueError:
-        await message.answer('Напиши юзернейм в формате "@username"')
-        return
 
-    kb = InlineKeyboardMarkup(
-        inline_keyboard=[
-            [
-                InlineKeyboardButton(
-                    text="Близкий друг",
-                    callback_data=callbacks.LinkCallback(
-                        username_to=username_to, rating=3
-                    ).pack(),
-                )
-            ],
-            [
-                InlineKeyboardButton(
-                    text="Приятель",
-                    callback_data=callbacks.LinkCallback(
-                        username_to=username_to, rating=2
-                    ).pack(),
-                ),
-            ],
-            [
-                InlineKeyboardButton(
-                    text="Знакомый",
-                    callback_data=callbacks.LinkCallback(
-                        username_to=username_to, rating=1
-                    ).pack(),
-                ),
-            ],
-        ]
-    )
+    main_user = await userdb.get_user(message.from_user.username)
+    if len(main_user.links) < 5:
+        await message.answer("Для доступа к реферальной программе отметь 5\\+ связей")
+        return
 
-    await message.answer("Кто он для тебя?", reply_markup=kb)
+    link = await create_start_link(bot, message.from_user.username, encode=True)
 
+    img = qrcode.make(link)
+    img_byte_arr = io.BytesIO()
+    img.save(img_byte_arr, format="PNG")
+    img_byte_arr = img_byte_arr.getvalue()
+    qr_file = BufferedInputFile(img_byte_arr, f"qr_{message.from_user.id}.png")
 
-@dp.callback_query(callbacks.LinkCallback.filter())
-async def process_data(query: CallbackQuery, callback_data: callbacks.LinkCallback):
-    from_username = query.from_user.username
-    if await userdb.get_user(from_username) is None:
-        await query.message.answer("Похоже тебе нужно перезапустить бота: /start")
+    if not main_user.invited:
+        await message.answer_photo(
+            photo=qr_file,
+            caption=generate_message(link),
+            parse_mode=ParseMode.MARKDOWN_V2,
+        )
         return
-    await userdb.add_link(
-        from_username,
-        Link(username_to=callback_data.username_to, rating=callback_data.rating),
+    users = await userdb.get_users(username=main_user.invited)
+    str_list = []
+    points = 0
+    for user in users:
+        str_list.append(
+            "• @" + user.username + " - " + ("🟡" if len(user.links) < 5 else "🟢")
+        )
+        if len(user.links) >= 5:
+            points += 1
+
+    await message.answer_photo(
+        photo=qr_file,
+        caption=generate_message(link, str_list, points),
+        parse_mode=ParseMode.MARKDOWN_V2,
     )
-    await query.message.edit_text(
-        as_list(
-            f"✅ @{callback_data.username_to} добавлен как {rating_to_text(callback_data.rating).lower()}",
-            "\n📝 Чтобы добавить ещё друга — просто введи следующий юзернейм.",
-            "\n🔁 Чем больше друзей ты добавишь — тем точнее будет твой социальный портрет!",
-        ),
-        reply_markup=None,
+    await message.answer(
+        "Давай посмотрим, кто перешёл по твоей ссылке\n"
+        + "\n".join(str_list)
+        + f"\nВсего баллов: {points}"
     )

src/models.py

@@ -1,11 +1,13 @@
 import re
-from typing import Annotated, Literal
+from typing import Annotated, Literal, Optional
 
 from pydantic import AfterValidator, BaseModel, Field
 
 
 def check_tg_username(username: str) -> str:
     username = username.strip().strip("@")
+    if len(username) < 5:
+        raise ValueError(f"{username} is not valid username")
     pattern = r"^[A-z0-9_]+$"
     if re.match(pattern, username):
         return username
@@ -32,11 +34,7 @@ class User(BaseModel):
     course: int = Field(ge=1, le=2)
     living: Living
 
-    _links: list[Link] = []
-
-    def set_link(self, link: Link):
-        for i in self._links:
-            if i.username_to == link.username_to:
-                i = link
-                return
-        self._links.append(link)
+    links: list[Link] = []
+
+    invited: list[Username] = []
+    invited_by: Optional[Username] = None