gigacook

1
import logging
2
from os import environ
3

4
from dotenv import load_dotenv
5
from telegram import BotCommand, InlineKeyboardButton, InlineKeyboardMarkup, MenuButtonCommands, Update
6
from telegram.ext import Application, CallbackQueryHandler, CommandHandler, ContextTypes, MessageHandler
7
from telegram.ext.filters import COMMAND, TEXT
8

9
from db import Bookmark, Product, Stage, run_async_session
10
from gpt import gpt_call
11

12
load_dotenv()
13

14
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO)
15
logging.getLogger("httpx").setLevel(logging.WARNING)
16

17
logger = logging.getLogger(__name__)
18

19

20
product_create_button = [[InlineKeyboardButton("Добавить ещё", callback_data="product_create")]]
21

22

23
async def command_start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
24
    commands = [
25
        ("start", "Начало"),
26
        ("products", "Товары"),
27
        ("recipe", "Получить рецепт"),
28
        ("bookmarks", "Закладки с рецептами"),
29
    ]
30

31
    await Stage.set(user=update.effective_user.id, name="start")
32

33
    await update.effective_user.get_bot().set_my_commands([BotCommand(c, d) for c, d in commands])
34
    await update.effective_user.set_menu_button(MenuButtonCommands())
35

36
    await update.message.reply_html(
37
        "Данный бот запоминает товары и генерирует рецепт по ним.\n"
38
        f"Список команд:\n"
39
        f"{'\n'.join([f'/{c} - {d}' for c, d in commands])}"
40
    )
41

42

43
async def command_products(update: Update, context: ContextTypes.DEFAULT_TYPE):
44
    await Stage.set(user=update.effective_user.id, name="product_list")
45

46
    buttons = [
47
        [InlineKeyboardButton(f"{product.name} [удалить]", callback_data=f"product_delete {product.id}")]
48
        for product in await Product.read(update.effective_user.id)
49
    ] + product_create_button
50

51
    await update.effective_user.send_message("Текущие продукты:", reply_markup=InlineKeyboardMarkup(buttons))
52

53

54
async def command_recipe(update: Update, context: ContextTypes.DEFAULT_TYPE):
55
    await Stage.set(user=update.effective_user.id, name="recipe")
56

57
    products = await Product.read(update.effective_user.id)
58
    if len(products) == 0:
59
        await update.effective_user.send_message("У вас нет товаров")
60

61
    payload = f"Напиши рецепт одного блюда используя только " f"{', '.join([p.name for p in products])}"
62

63
    await update.effective_user.send_message(
64
        gpt_call(payload),
65
        reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Сохранить", callback_data="bookmark_create")]]),
66
    )
67

68

69
async def command_bookmarks(update: Update, context: ContextTypes.DEFAULT_TYPE):
70
    await Stage.set(user=update.effective_user.id, name="bookmark_list")
71

72
    for bookmark in await Bookmark.read(update.effective_user.id):
73
        await update.effective_user.send_message(
74
            bookmark.text,
75
            reply_markup=InlineKeyboardMarkup(
76
                [[InlineKeyboardButton("удалить", callback_data=f"bookmarks_delete {bookmark.id}")]]
77
            ),
78
        )
79
    else:
80
        await update.effective_user.send_message("Закладок нет")
81

82

83
async def message(update: Update, context: ContextTypes.DEFAULT_TYPE):
84
    text = update.message.text
85
    stage = (await Stage.read(update.effective_user.id))[0].name
86

87
    if stage == "product_create":
88
        await Product.create(user=update.effective_user.id, name=text)
89
        await command_products(update, context)
90
    else:
91
        await update.effective_user.send_message("Сообщение не распознано")
92

93

94
async def button(update: Update, context: ContextTypes.DEFAULT_TYPE):
95
    await update.callback_query.answer()
96

97
    arr = update.callback_query.data.split()
98
    data, pk = arr[0], int(arr[1]) if len(arr) == 2 else None
99

100
    await Stage.set(user=update.effective_user.id, name=data)
101

102
    if data == "product_create":
103
        await update.effective_user.send_message("Введите название продукта")
104

105
    elif data == "product_delete":
106
        await Product.delete(pk)
107
        await update.effective_user.send_message("Продукт удалён")
108
        await command_products(update, context)
109

110
    elif data == "bookmark_create":
111
        await Bookmark.create(user=update.effective_user.id, text=update.callback_query.message.text)
112
        await update.effective_user.send_message("Рецепт сохранён в закладки")
113

114
    elif data == "bookmarks_delete":
115
        await Bookmark.delete(pk)
116
        await update.effective_user.send_message("Продукт удалён")
117
        await command_bookmarks(update, context)
118

119

120
def run_polling_tg_app():
121
    application = Application.builder().token(environ["TELEGRAM_TOKEN"]).build()
122

123
    application.add_handler(CommandHandler("start", command_start))
124
    application.add_handler(CommandHandler("products", command_products))
125
    application.add_handler(CommandHandler("recipe", command_recipe))
126
    application.add_handler(CommandHandler("bookmarks", command_bookmarks))
127

128
    application.add_handler(MessageHandler(TEXT & ~COMMAND, message))
129

130
    application.add_handler(CallbackQueryHandler(button))
131

132
    application.run_polling(allowed_updates=Update.ALL_TYPES)
133

134

135
if __name__ == "__main__":
136
    run_async_session()
137

138
    run_polling_tg_app()
139