diff --git a/pyproject.toml b/pyproject.toml index d0c152f..1cc199a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,4 +10,5 @@ dependencies = [ "dotenv>=0.9.9", "pydantic>=2.11.10", "pymongo>=4.15.4", + "qrcode[pil]>=8.2", ] diff --git a/requirements.txt b/requirements.txt index 98f0af3..193b373 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file +pymongo>=4.15.4 +pymongo>=4.15.4 +qrcode[pil]>=8.2 \ No newline at end of file diff --git a/src/bot.py b/src/bot.py index b1802cc..f06415b 100644 --- a/src/bot.py +++ b/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}" ) diff --git a/src/models.py b/src/models.py index 8dda885..40fe1b6 100644 --- a/src/models.py +++ b/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 diff --git a/src/userdb.py b/src/userdb.py index 7fa169c..68c7ef3 100644 --- a/src/userdb.py +++ b/src/userdb.py @@ -1,13 +1,10 @@ import asyncio -import os -from typing import List, Optional +from typing import Iterable, List, Optional +from aiogram.types import user import pymongo.asynchronous.collection as pymongo_collection import pymongo.asynchronous.database as pymongo_database -from dotenv import load_dotenv -from pymongo import AsyncMongoClient from pymongo.asynchronous.cursor import AsyncCursor -from pymongo.errors import ServerSelectionTimeoutError from .models import Link, User, Username @@ -23,7 +20,11 @@ class UserDB: def __init__(self, client): self.db = client.get_database("cu_graph_bot") self.collection = self.db.get_collection("users") - asyncio.create_task(self.collection.create_index("username", unique=True)) + asyncio.gather( + self.collection.create_index("username", unique=True), + # NOTE: Also temporary + self.collection.update_many({}, {"$rename": {"_links": "links"}}), + ) async def add_user(self, user: User) -> None: current_user_to_add = await self.collection.find_one( @@ -38,34 +39,55 @@ async def get_links(self, username: Username) -> List[Link]: all_friends = await self.collection.find_one({"username": username}) if all_friends is None: raise UserNotExist - if "_links" not in all_friends: + if "links" not in all_friends: return [] - return [Link.model_validate(i) for i in all_friends["_links"]] + links_list = [] + pop_links = [] + for i in all_friends["links"]: + try: + links_list.append(Link.model_validate(i)) + except: + pop_links.append(i["username_to"]) + await self.collection.update_one( + {"username": username}, + {"$pull": {"links": {"username_to": {"$in": pop_links}}}}, + ) + return links_list async def get_user(self, username: Username) -> Optional[User]: - all_user_data = await self.collection.find_one({"username": username}) - return all_user_data - - async def get_users(self, links_less_than: int) -> AsyncCursor: - all_user_data = self.collection.find( - {f"_links.{links_less_than}": {"$exists": False}} - ) + user_data = await self.collection.find_one({"username": username}) + return User.model_validate(user_data) if user_data is not None else None + + async def get_users( + self, + username: Username | Iterable[Username] | None = None, + links_less_than: Optional[int] = None, + ) -> AsyncCursor: + query = {} + if isinstance(username, (list, tuple)): + query["$or"] = [{"username": i} for i in username] + elif username: + query["username"] = username + if links_less_than: + query[f"links.{links_less_than}"] = {"$exists": False} + all_user_data = self.collection.find(query) + all_user_data = [User.model_validate(i) async for i in all_user_data] return all_user_data async def add_link(self, username: Username, link: Link) -> None: current_user = await self.collection.find_one({"username": username}) if current_user is None: raise UserNotExist - for ind, ilink in enumerate(current_user.get("_links", [])): + for ind, ilink in enumerate(current_user.get("links", [])): if ilink["username_to"] == link.username_to: await self.collection.update_one( {"username": username}, - {"$set": {f"_links.{ind}.rating": link.rating}}, + {"$set": {f"links.{ind}.rating": link.rating}}, ) return else: await self.collection.update_one( - {"username": username}, {"$push": {"_links": link.model_dump()}} + {"username": username}, {"$push": {"links": link.model_dump()}} ) async def count_users(self, links_more_than: int = 4) -> int: @@ -73,7 +95,7 @@ async def count_users(self, links_more_than: int = 4) -> int: [ i async for i in self.collection.find( - {f"_links.{links_more_than}": {"$exists": True}} + {f"links.{links_more_than}": {"$exists": True}} ) ] ) @@ -81,9 +103,21 @@ async def count_users(self, links_more_than: int = 4) -> int: # NOTE: This is temporary function for fixing current database # And should be deleted someday because of obvious performance loss - async def add_ids_to_user(self, username: str, userid: int, chatid: int) -> None: + async def add_ids_to_user( + self, username: Username, userid: int, chatid: int + ) -> None: await self.collection.update_one( {"username": username}, {"$set": {"userid": userid, "chatid": chatid}} ) + async def add_invited(self, username: Username, username_invited: Username) -> None: + await self.collection.update_one( + {"username": username}, {"$addToSet": {"invited": username_invited}} + ) + async def add_invited_by( + self, username: Username, username_invited_by: Username + ) -> None: + await self.collection.update_one( + {"username": username}, {"$set": {"invited_by": username_invited_by}} + )