scheduleSon/backend/telegram_bot.py

493 lines
20 KiB
Python
Raw Normal View History

from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (
Application, CommandHandler, CallbackQueryHandler, MessageHandler,
ContextTypes, ConversationHandler, filters
)
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from datetime import datetime, timedelta
import pytz
import os
import asyncio
import logging
from backend.database import AsyncSessionLocal, Task, Event, TelegramUser
from backend.utils import get_week_range, format_date_russian, check_event_overlap, get_weekday_from_date
from backend.models import TaskCreate, EventCreate
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
TZ = pytz.timezone("Europe/London")
BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
# Состояния для ConversationHandler
SELECTING_DATE, SELECTING_HOUR, SELECTING_MINUTE, SELECTING_DURATION_HOUR, SELECTING_DURATION_MIN, ENTERING_TITLE = range(6)
SELECTING_TASK_DATE, ENTERING_TASK_TITLE, SELECTING_REPEAT, SELECTING_COPY_DAYS = range(10, 14)
user_states = {} # Временное хранилище состояний пользователей
async def register_user(chat_id: str):
"""Зарегистрировать пользователя для получения напоминаний"""
async with AsyncSessionLocal() as db:
result = await db.execute(
select(TelegramUser).where(TelegramUser.chat_id == chat_id)
)
user = result.scalar_one_or_none()
if not user:
user = TelegramUser(chat_id=chat_id)
db.add(user)
await db.commit()
async def get_schedule_for_date(date_str: str) -> str:
"""Получить расписание на дату в текстовом формате"""
async with AsyncSessionLocal() as db:
# Получаем tasks (разовые)
result = await db.execute(
select(Task).where(
and_(Task.date == date_str, Task.repeat_weekly == False)
)
)
tasks = result.scalars().all()
# Получаем weekly tasks для этой даты
from backend.utils import materialize_weekly_tasks
weekly_items = await materialize_weekly_tasks(db, date_str, date_str)
weekly_tasks = [item for item in weekly_items if item.kind == "task"]
# Получаем events
result = await db.execute(
select(Event).where(Event.date == date_str)
)
events = result.scalars().all()
events = sorted(events, key=lambda e: e.start_time)
text = f"📅 {format_date_russian(date_str)}\n\n"
# Объединяем tasks
all_tasks = []
for task in tasks:
all_tasks.append(task.title)
for task_item in weekly_tasks:
all_tasks.append(task_item.title)
if all_tasks:
text += "📋 Задачи:\n"
for title in all_tasks:
text += f"{title}\n"
text += "\n"
if events:
text += "⏰ Занятия:\n"
for event in events:
end_time = calculate_end_time(event.start_time, event.duration_min)
text += f" {event.start_time}-{end_time} {event.title}\n"
if not all_tasks and not events:
text += "На этот день расписания нет.\n"
return text
def calculate_end_time(start_time: str, duration_min: int) -> str:
"""Вычислить время окончания"""
hour, minute = map(int, start_time.split(":"))
total_minutes = hour * 60 + minute + duration_min
end_hour = total_minutes // 60
end_minute = total_minutes % 60
return f"{end_hour:02d}:{end_minute:02d}"
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Обработчик команды /start"""
await register_user(str(update.effective_chat.id))
keyboard = [
[InlineKeyboardButton("📅 Сегодня", callback_data="schedule_today")],
[InlineKeyboardButton("📅 Завтра", callback_data="schedule_tomorrow")],
[InlineKeyboardButton("📅 Оставшиеся дни недели", callback_data="schedule_week")],
[InlineKeyboardButton(" Добавить занятие", callback_data="add_event")],
[InlineKeyboardButton(" Добавить задачу", callback_data="add_task")],
]
reply_markup = InlineKeyboardMarkup(keyboard)
await update.message.reply_text(
"Выберите действие:",
reply_markup=reply_markup
)
async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Обработчик кнопок"""
query = update.callback_query
await query.answer()
chat_id = str(query.from_user.id)
today = datetime.now(TZ).date()
if query.data == "schedule_today":
date_str = today.strftime("%Y-%m-%d")
text = await get_schedule_for_date(date_str)
await query.edit_message_text(text)
elif query.data == "schedule_tomorrow":
tomorrow = today + timedelta(days=1)
date_str = tomorrow.strftime("%Y-%m-%d")
text = await get_schedule_for_date(date_str)
await query.edit_message_text(text)
elif query.data == "schedule_week":
# Оставшиеся дни недели
weekday = today.weekday()
text = ""
for i in range(weekday, 7):
date = today + timedelta(days=i - weekday)
date_str = date.strftime("%Y-%m-%d")
day_text = await get_schedule_for_date(date_str)
text += day_text + "\n"
await query.edit_message_text(text[:4000]) # Telegram limit
elif query.data == "add_event":
# Начинаем процесс добавления event
user_states[chat_id] = {"action": "add_event", "step": "date"}
await show_date_selection(query, "event")
return SELECTING_DATE
elif query.data == "add_task":
# Начинаем процесс добавления task
user_states[chat_id] = {"action": "add_task", "step": "date"}
await show_date_selection(query, "task")
return SELECTING_TASK_DATE
elif query.data.startswith("date_"):
# Выбрана дата
date_str = query.data.replace("date_", "")
if chat_id not in user_states:
user_states[chat_id] = {}
user_states[chat_id]["date"] = date_str
if user_states[chat_id].get("action") == "add_event":
await show_hour_selection(query)
return SELECTING_HOUR
else: # add_task
await query.edit_message_text("Введите название задачи:")
return ENTERING_TASK_TITLE
elif query.data.startswith("hour_"):
hour = query.data.replace("hour_", "")
user_states[chat_id]["hour"] = hour
await show_minute_selection(query)
return SELECTING_MINUTE
elif query.data.startswith("minute_"):
minute = query.data.replace("minute_", "")
user_states[chat_id]["minute"] = minute
await show_duration_selection(query)
return SELECTING_DURATION_HOUR
elif query.data.startswith("dur_hour_"):
dur_hour = int(query.data.replace("dur_hour_", ""))
user_states[chat_id]["duration_hour"] = dur_hour
await show_duration_minute_selection(query)
return SELECTING_DURATION_MIN
elif query.data.startswith("dur_min_"):
dur_min = int(query.data.replace("dur_min_", ""))
user_states[chat_id]["duration_minute"] = dur_min
await query.edit_message_text("Введите название занятия:")
return ENTERING_TITLE
# Возвращаемся в главное меню
keyboard = [
[InlineKeyboardButton("📅 Сегодня", callback_data="schedule_today")],
[InlineKeyboardButton("📅 Завтра", callback_data="schedule_tomorrow")],
[InlineKeyboardButton("📅 Оставшиеся дни недели", callback_data="schedule_week")],
[InlineKeyboardButton(" Добавить занятие", callback_data="add_event")],
[InlineKeyboardButton(" Добавить задачу", callback_data="add_task")],
]
reply_markup = InlineKeyboardMarkup(keyboard)
await query.edit_message_text("Выберите действие:", reply_markup=reply_markup)
return ConversationHandler.END
async def show_date_selection(query, kind="event"):
"""Показать выбор даты"""
today = datetime.now(TZ).date()
monday = today - timedelta(days=today.weekday())
next_monday = monday + timedelta(days=7)
keyboard = []
weekdays = ["Пн", "Вт", "Ср", "Чт", "Пт", "Сб", "Вс"]
# Текущая неделя
for i in range(7):
date = monday + timedelta(days=i)
date_str = date.strftime("%Y-%m-%d")
label = f"{weekdays[i]} {date.day}"
keyboard.append([InlineKeyboardButton(label, callback_data=f"date_{date_str}")])
# Следующая неделя
for i in range(7):
date = next_monday + timedelta(days=i)
date_str = date.strftime("%Y-%m-%d")
label = f"{weekdays[i]} {date.day}"
keyboard.append([InlineKeyboardButton(label, callback_data=f"date_{date_str}")])
reply_markup = InlineKeyboardMarkup(keyboard)
await query.edit_message_text("Выберите дату:", reply_markup=reply_markup)
async def show_hour_selection(query):
"""Показать выбор часа"""
keyboard = []
for hour in range(8, 21):
keyboard.append([InlineKeyboardButton(f"{hour:02d}:00", callback_data=f"hour_{hour:02d}")])
reply_markup = InlineKeyboardMarkup(keyboard)
await query.edit_message_text("Выберите час начала:", reply_markup=reply_markup)
async def show_minute_selection(query):
"""Показать выбор минут"""
keyboard = [
[InlineKeyboardButton("00", callback_data="minute_00")],
[InlineKeyboardButton("15", callback_data="minute_15")],
[InlineKeyboardButton("30", callback_data="minute_30")],
[InlineKeyboardButton("45", callback_data="minute_45")],
]
reply_markup = InlineKeyboardMarkup(keyboard)
await query.edit_message_text("Выберите минуты:", reply_markup=reply_markup)
async def show_duration_selection(query):
"""Показать выбор длительности (часы)"""
keyboard = []
for hour in range(0, 4):
keyboard.append([InlineKeyboardButton(f"{hour} ч", callback_data=f"dur_hour_{hour}")])
reply_markup = InlineKeyboardMarkup(keyboard)
await query.edit_message_text("Выберите длительность (часы):", reply_markup=reply_markup)
async def show_duration_minute_selection(query):
"""Показать выбор длительности (минуты)"""
keyboard = [
[InlineKeyboardButton("00", callback_data="dur_min_00")],
[InlineKeyboardButton("15", callback_data="dur_min_15")],
[InlineKeyboardButton("30", callback_data="dur_min_30")],
[InlineKeyboardButton("45", callback_data="dur_min_45")],
]
reply_markup = InlineKeyboardMarkup(keyboard)
await query.edit_message_text("Выберите длительность (минуты):", reply_markup=reply_markup)
async def handle_event_title(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Обработка ввода названия event"""
chat_id = str(update.effective_user.id)
title = update.message.text
state = user_states.get(chat_id, {})
if not state:
await update.message.reply_text("Ошибка: сессия истекла. Начните заново.")
return ConversationHandler.END
date = state.get("date")
hour = state.get("hour")
minute = state.get("minute")
dur_hour = state.get("duration_hour", 0)
dur_min = state.get("duration_minute", 0)
if not all([date, hour, minute]):
await update.message.reply_text("Ошибка: не все данные заполнены. Начните заново.")
if chat_id in user_states:
del user_states[chat_id]
return ConversationHandler.END
duration_min = dur_hour * 60 + dur_min
if duration_min == 0:
duration_min = 15 # Минимум 15 минут
start_time = f"{hour}:{minute}"
# Создаём event
async with AsyncSessionLocal() as db:
# Проверка пересечений
overlap = await check_event_overlap(db, date, start_time, duration_min, None)
if overlap:
await update.message.reply_text("❌ Нельзя добавить: пересечение по времени")
if chat_id in user_states:
del user_states[chat_id]
return ConversationHandler.END
event = Event(
date=date,
start_time=start_time,
duration_min=duration_min,
title=title
)
db.add(event)
await db.commit()
await update.message.reply_text(f"✅ Занятие добавлено: {date} {start_time} - {title}")
if chat_id in user_states:
del user_states[chat_id]
return ConversationHandler.END
async def handle_task_title(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Обработка ввода названия task"""
chat_id = str(update.effective_user.id)
title = update.message.text
state = user_states.get(chat_id, {})
state["title"] = title
keyboard = [
[InlineKeyboardButton("Да, повторять каждую неделю", callback_data="repeat_yes")],
[InlineKeyboardButton("Нет, только на эту дату", callback_data="repeat_no")],
]
reply_markup = InlineKeyboardMarkup(keyboard)
await update.message.reply_text(
"Повторять эту задачу каждую неделю?",
reply_markup=reply_markup
)
return SELECTING_REPEAT
async def handle_repeat_choice(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Обработка выбора повторения"""
query = update.callback_query
await query.answer()
chat_id = str(query.from_user.id)
state = user_states.get(chat_id, {})
if not state:
await query.edit_message_text("Ошибка: сессия истекла. Начните заново.")
return ConversationHandler.END
if query.data == "repeat_yes":
state["repeat_weekly"] = True
# Создаём weekly task
async with AsyncSessionLocal() as db:
weekday = get_weekday_from_date(state["date"])
task = Task(
date=state["date"],
title=state["title"],
repeat_weekly=True,
weekday=weekday
)
db.add(task)
await db.commit()
await query.edit_message_text(f"✅ Задача добавлена (повторяется каждую неделю): {state['title']}")
else:
state["repeat_weekly"] = False
# Создаём разовую task
async with AsyncSessionLocal() as db:
task = Task(
date=state["date"],
title=state["title"],
repeat_weekly=False
)
db.add(task)
await db.commit()
await query.edit_message_text(f"✅ Задача добавлена: {state['title']}")
del user_states[chat_id]
return ConversationHandler.END
async def send_reminders():
"""Отправка напоминаний о событиях"""
while True:
try:
now = datetime.now(TZ)
target_time = now + timedelta(minutes=5)
target_date = target_time.date()
target_time_str = target_time.strftime("%H:%M")
async with AsyncSessionLocal() as db:
# Находим события, которые начинаются через 5 минут
result = await db.execute(
select(Event).where(
and_(
Event.date == target_date.strftime("%Y-%m-%d"),
Event.start_time == target_time_str
)
)
)
events = result.scalars().all()
if events:
# Получаем всех пользователей
result = await db.execute(select(TelegramUser))
users = result.scalars().all()
for event in events:
message = f"🔔 Напоминание: через 5 минут начинается {event.title} ({event.start_time})"
for user in users:
try:
await application.bot.send_message(
chat_id=user.chat_id,
text=message
)
except Exception as e:
logger.error(f"Error sending reminder to {user.chat_id}: {e}")
except Exception as e:
logger.error(f"Error in reminder loop: {e}")
await asyncio.sleep(60) # Проверяем каждую минуту
application = None
async def start_bot():
"""Запустить бота"""
global application
if not BOT_TOKEN:
logger.warning("TELEGRAM_BOT_TOKEN not set, bot will not start")
return
application = Application.builder().token(BOT_TOKEN).build()
# Conversation handler для добавления event
event_conv = ConversationHandler(
entry_points=[CallbackQueryHandler(button_handler, pattern="^add_event$")],
states={
SELECTING_DATE: [CallbackQueryHandler(button_handler, pattern="^date_")],
SELECTING_HOUR: [CallbackQueryHandler(button_handler, pattern="^hour_")],
SELECTING_MINUTE: [CallbackQueryHandler(button_handler, pattern="^minute_")],
SELECTING_DURATION_HOUR: [CallbackQueryHandler(button_handler, pattern="^dur_hour_")],
SELECTING_DURATION_MIN: [CallbackQueryHandler(button_handler, pattern="^dur_min_")],
ENTERING_TITLE: [MessageHandler(filters.TEXT & ~filters.COMMAND, handle_event_title)],
},
fallbacks=[],
name="event_conv",
)
# Conversation handler для добавления task
task_conv = ConversationHandler(
entry_points=[CallbackQueryHandler(button_handler, pattern="^add_task$")],
states={
SELECTING_TASK_DATE: [CallbackQueryHandler(button_handler, pattern="^date_")],
ENTERING_TASK_TITLE: [MessageHandler(filters.TEXT & ~filters.COMMAND, handle_task_title)],
SELECTING_REPEAT: [CallbackQueryHandler(handle_repeat_choice, pattern="^repeat_")],
},
fallbacks=[],
name="task_conv",
)
application.add_handler(CommandHandler("start", start))
application.add_handler(event_conv)
application.add_handler(task_conv)
application.add_handler(CallbackQueryHandler(button_handler))
# Запуск напоминаний
asyncio.create_task(send_reminders())
# Запускаем бота в фоне
await application.initialize()
await application.start()
await application.updater.start_polling(drop_pending_updates=True)