366 lines
16 KiB
Python
366 lines
16 KiB
Python
"""Async handlers for private messages."""
|
||
import re
|
||
from typing import Optional, Dict
|
||
from telebot.async_telebot import AsyncTeleBot
|
||
from telebot import types
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy import select
|
||
from bot.database import get_db_session, User, Chat, UserChat
|
||
from bot.messages import THEMES, get_theme_emoji
|
||
from bot.logger import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
# User states for conversation flow
|
||
user_states: Dict[int, str] = {}
|
||
|
||
|
||
def register_private_handlers(bot: AsyncTeleBot) -> None:
|
||
"""Register all private message handlers."""
|
||
|
||
@bot.message_handler(commands=['start'], chat_types=['private'])
|
||
async def handle_start(message: types.Message) -> None:
|
||
"""Handle /start command in private chat."""
|
||
if not message.from_user:
|
||
return
|
||
|
||
user_id = message.from_user.id
|
||
username: Optional[str] = message.from_user.username
|
||
first_name: str = message.from_user.first_name or "Пользователь"
|
||
logger.info(f"User {user_id} ({first_name}) started conversation")
|
||
|
||
async for db in get_db_session():
|
||
try:
|
||
# Check if user exists
|
||
result = await db.execute(select(User).filter(User.user_id == user_id))
|
||
user = result.scalar_one_or_none()
|
||
|
||
if user:
|
||
# User already exists - ask if they want to update
|
||
logger.debug(f"User {user_id} already exists in database")
|
||
await bot.send_message(
|
||
user_id,
|
||
f"Привет, {first_name}! Ты уже поделился со мной днем рождения.\n"
|
||
f"Используй /update, чтобы обновить свои данные."
|
||
)
|
||
else:
|
||
# New user - ask for birthday
|
||
logger.info(f"New user {user_id} ({first_name}) - requesting birthday")
|
||
user_states[user_id] = 'waiting_birthday'
|
||
await bot.send_message(
|
||
user_id,
|
||
f"Привет, {first_name}! 👋\n\n"
|
||
f"Пришли мне свой день рождения в формате ДД.ММ или ДД.ММ.ГГГГ\n"
|
||
f"Например: 15.03 или 15.03.1990"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Error handling /start for user {user_id}: {e}", exc_info=True)
|
||
finally:
|
||
await db.close()
|
||
|
||
@bot.message_handler(commands=['update'], chat_types=['private'])
|
||
async def handle_update(message: types.Message) -> None:
|
||
"""Handle /update command in private chat."""
|
||
if not message.from_user:
|
||
return
|
||
|
||
user_id = message.from_user.id
|
||
first_name: str = message.from_user.first_name or "Пользователь"
|
||
logger.info(f"User {user_id} ({first_name}) requested update")
|
||
|
||
async for db in get_db_session():
|
||
try:
|
||
result = await db.execute(select(User).filter(User.user_id == user_id))
|
||
user = result.scalar_one_or_none()
|
||
if user:
|
||
user_states[user_id] = 'waiting_birthday'
|
||
await bot.send_message(
|
||
user_id,
|
||
f"Хорошо, {first_name}! Пришли мне свой день рождения в формате ДД.ММ или ДД.ММ.ГГГГ"
|
||
)
|
||
else:
|
||
await bot.send_message(
|
||
user_id,
|
||
"Ты еще не поделился со мной днем рождения. Используй /start для начала."
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Error handling /update for user {user_id}: {e}", exc_info=True)
|
||
finally:
|
||
await db.close()
|
||
|
||
@bot.message_handler(commands=['help'], chat_types=['private'])
|
||
async def handle_help(message: types.Message) -> None:
|
||
"""Handle /help command in private chat."""
|
||
help_text = (
|
||
"📚 Команды бота:\n\n"
|
||
"/start - Поделиться днем рождения и выбрать тему предпочтений\n"
|
||
"/update - Обновить свой день рождения или тему предпочтений\n"
|
||
"/help - Показать это сообщение\n\n"
|
||
"В группах доступны команды:\n"
|
||
"/stats - Статистика по чату\n"
|
||
"/week - Дни рождения на 7 дней\n"
|
||
"/month - Дни рождения на 31 день\n\n"
|
||
"from olly & cursor with <3"
|
||
)
|
||
await bot.send_message(message.chat.id, help_text)
|
||
|
||
@bot.message_handler(func=lambda m: m.chat.type == 'private' and m.from_user and m.from_user.id in user_states and m.text)
|
||
async def handle_birthday_input(message: types.Message) -> None:
|
||
"""Handle birthday input from user."""
|
||
if not message.from_user or not message.text:
|
||
return
|
||
|
||
user_id = message.from_user.id
|
||
state = user_states.get(user_id)
|
||
|
||
if state == 'waiting_birthday':
|
||
# Parse birthday
|
||
birthday_text = message.text.strip()
|
||
parsed = parse_birthday(birthday_text)
|
||
|
||
if not parsed:
|
||
await bot.send_message(
|
||
user_id,
|
||
"Неверный формат даты. Используй формат ДД.ММ или ДД.ММ.ГГГГ\n"
|
||
"Например: 15.03 или 15.03.1990"
|
||
)
|
||
return
|
||
|
||
day, month, year = parsed
|
||
|
||
# Validate date
|
||
if not is_valid_date(day, month, year):
|
||
await bot.send_message(
|
||
user_id,
|
||
"Неверная дата. Проверь правильность дня и месяца."
|
||
)
|
||
return
|
||
|
||
# Save birthday and ask for preference
|
||
async for db in get_db_session():
|
||
try:
|
||
if not message.from_user:
|
||
return
|
||
username: Optional[str] = message.from_user.username
|
||
first_name: str = message.from_user.first_name or "Пользователь"
|
||
|
||
result = await db.execute(select(User).filter(User.user_id == user_id))
|
||
user = result.scalar_one_or_none()
|
||
if user:
|
||
# Update existing user
|
||
logger.info(f"Updating birthday for user {user_id}: {day:02d}.{month:02d}.{year or 'N/A'}")
|
||
user.birthday_day = day
|
||
user.birthday_month = month
|
||
user.birthday_year = year
|
||
user.first_name = first_name
|
||
user.username = username
|
||
else:
|
||
# Create new user
|
||
logger.info(f"Creating new user {user_id} with birthday: {day:02d}.{month:02d}.{year or 'N/A'}")
|
||
user = User(
|
||
user_id=user_id,
|
||
username=username,
|
||
first_name=first_name,
|
||
birthday_day=day,
|
||
birthday_month=month,
|
||
birthday_year=year,
|
||
preference_theme="Музыка" # Default theme
|
||
)
|
||
db.add(user)
|
||
|
||
await db.commit()
|
||
|
||
# Ask for preference theme
|
||
user_states[user_id] = 'waiting_preference'
|
||
await ask_preference_theme(bot, user_id)
|
||
except Exception as e:
|
||
logger.error(f"Error saving birthday for user {user_id}: {e}", exc_info=True)
|
||
finally:
|
||
await db.close()
|
||
|
||
elif state == 'waiting_preference':
|
||
# User selected preference theme
|
||
if not message.text:
|
||
return
|
||
|
||
theme_text = message.text.strip()
|
||
|
||
if theme_text not in THEMES:
|
||
await bot.send_message(
|
||
user_id,
|
||
"Пожалуйста, выбери одну из предложенных тем, нажав на кнопку."
|
||
)
|
||
return
|
||
|
||
# Save preference
|
||
async for db in get_db_session():
|
||
try:
|
||
result = await db.execute(select(User).filter(User.user_id == user_id))
|
||
user = result.scalar_one_or_none()
|
||
if user:
|
||
logger.info(f"User {user_id} selected theme: {theme_text}")
|
||
user.preference_theme = theme_text
|
||
await db.commit()
|
||
|
||
# Update user in all chats
|
||
await update_user_in_chats(bot, db, user)
|
||
|
||
await bot.send_message(
|
||
user_id,
|
||
f"Отлично! Я запомнил твои предпочтения: {theme_text}\n\n"
|
||
f"Теперь я буду поздравлять тебя с днем рождения во всех чатах, где ты состоишь!"
|
||
)
|
||
user_states.pop(user_id, None)
|
||
except Exception as e:
|
||
logger.error(f"Error saving preference for user {user_id}: {e}", exc_info=True)
|
||
finally:
|
||
await db.close()
|
||
|
||
@bot.callback_query_handler(func=lambda call: call.data and call.data.startswith('theme_'))
|
||
async def handle_theme_selection(call: types.CallbackQuery) -> None:
|
||
"""Handle theme selection from inline keyboard."""
|
||
if not call.from_user or not call.data or not call.message:
|
||
return
|
||
|
||
user_id = call.from_user.id
|
||
|
||
if user_states.get(user_id) != 'waiting_preference':
|
||
await bot.answer_callback_query(call.id, "Это действие недоступно сейчас.")
|
||
return
|
||
|
||
theme = call.data.replace('theme_', '')
|
||
|
||
if theme not in THEMES:
|
||
await bot.answer_callback_query(call.id, "Неверная тема.")
|
||
return
|
||
|
||
# Save preference
|
||
async for db in get_db_session():
|
||
try:
|
||
result = await db.execute(select(User).filter(User.user_id == user_id))
|
||
user = result.scalar_one_or_none()
|
||
if user:
|
||
logger.info(f"User {user_id} selected theme via callback: {theme}")
|
||
user.preference_theme = theme
|
||
await db.commit()
|
||
|
||
# Update user in all chats
|
||
await update_user_in_chats(bot, db, user)
|
||
|
||
await bot.answer_callback_query(call.id, f"Выбрано: {theme}")
|
||
await bot.edit_message_text(
|
||
f"Отлично! Я запомнил твои предпочтения: {theme}\n\n"
|
||
f"Теперь я буду поздравлять тебя с днем рождения во всех чатах, где ты состоишь!",
|
||
chat_id=call.message.chat.id,
|
||
message_id=call.message.message_id
|
||
)
|
||
user_states.pop(user_id, None)
|
||
except Exception as e:
|
||
logger.error(f"Error saving preference via callback for user {user_id}: {e}", exc_info=True)
|
||
finally:
|
||
await db.close()
|
||
|
||
|
||
def parse_birthday(text: str) -> Optional[tuple[int, int, Optional[int]]]:
|
||
"""Parse birthday from text. Returns (day, month, year) or None."""
|
||
# Try DD.MM.YYYY format
|
||
match = re.match(r'^(\d{1,2})\.(\d{1,2})\.(\d{4})$', text)
|
||
if match:
|
||
day, month, year = int(match.group(1)), int(match.group(2)), int(match.group(3))
|
||
return (day, month, year)
|
||
|
||
# Try DD.MM format
|
||
match = re.match(r'^(\d{1,2})\.(\d{1,2})$', text)
|
||
if match:
|
||
day, month = int(match.group(1)), int(match.group(2))
|
||
return (day, month, None)
|
||
|
||
return None
|
||
|
||
|
||
def is_valid_date(day: int, month: int, year: Optional[int]) -> bool:
|
||
"""Check if date is valid."""
|
||
if month < 1 or month > 12:
|
||
return False
|
||
if day < 1 or day > 31:
|
||
return False
|
||
|
||
# Check specific month limits
|
||
if month in [4, 6, 9, 11] and day > 30:
|
||
return False
|
||
if month == 2:
|
||
if year:
|
||
# Check leap year
|
||
if (year % 4 == 0 and year % 100 != 0) or (year % 400 == 0):
|
||
max_day = 29
|
||
else:
|
||
max_day = 28
|
||
else:
|
||
max_day = 29 # Assume leap year if year not provided
|
||
if day > max_day:
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
async def ask_preference_theme(bot: AsyncTeleBot, user_id: int) -> None:
|
||
"""Ask user to select preference theme."""
|
||
keyboard = types.InlineKeyboardMarkup(row_width=2)
|
||
|
||
# Show hobbies (themes) in 2 buttons per row
|
||
for i in range(0, len(THEMES), 2):
|
||
theme1 = THEMES[i]
|
||
emoji1 = get_theme_emoji(theme1)
|
||
btn1 = types.InlineKeyboardButton(
|
||
text=f"{emoji1} {theme1}",
|
||
callback_data=f'theme_{theme1}'
|
||
)
|
||
|
||
# Optional second button in the same row
|
||
if i + 1 < len(THEMES):
|
||
theme2 = THEMES[i + 1]
|
||
emoji2 = get_theme_emoji(theme2)
|
||
btn2 = types.InlineKeyboardButton(
|
||
text=f"{emoji2} {theme2}",
|
||
callback_data=f'theme_{theme2}'
|
||
)
|
||
keyboard.add(btn1, btn2)
|
||
else:
|
||
keyboard.add(btn1)
|
||
|
||
await bot.send_message(
|
||
user_id,
|
||
"Что тебе нравится? Выбери одну из тем:",
|
||
reply_markup=keyboard
|
||
)
|
||
|
||
|
||
async def update_user_in_chats(bot: AsyncTeleBot, db: AsyncSession, user: User) -> None:
|
||
"""Update user information in all chats where bot is admin."""
|
||
# Get all chats where bot is admin
|
||
result = await db.execute(select(Chat).filter(Chat.bot_is_admin == True))
|
||
admin_chats = result.scalars().all()
|
||
|
||
for chat in admin_chats:
|
||
try:
|
||
# Check if user is member of this chat
|
||
member = await bot.get_chat_member(chat.chat_id, user.user_id)
|
||
if member.status not in ['left', 'kicked']:
|
||
# User is in chat - add/update relationship
|
||
user_chat_result = await db.execute(
|
||
select(UserChat).filter(
|
||
UserChat.user_id == user.user_id,
|
||
UserChat.chat_id == chat.chat_id
|
||
)
|
||
)
|
||
user_chat = user_chat_result.scalar_one_or_none()
|
||
|
||
if not user_chat:
|
||
user_chat = UserChat(user_id=user.user_id, chat_id=chat.chat_id)
|
||
db.add(user_chat)
|
||
await db.commit()
|
||
except Exception:
|
||
# User might have blocked bot or bot was removed from chat
|
||
pass
|