forked from HSE_team/BetterCallPraskovia
382 lines
15 KiB
Python
382 lines
15 KiB
Python
"""
|
||
Обработчики для работы с документами
|
||
"""
|
||
from aiogram import Router, F
|
||
from aiogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton, CallbackQuery
|
||
from aiogram.filters import StateFilter
|
||
from aiogram.fsm.context import FSMContext
|
||
import aiohttp
|
||
from tg_bot.config.settings import settings
|
||
from tg_bot.infrastructure.telegram.states.collection_states import (
|
||
DocumentEditStates,
|
||
DocumentUploadStates
|
||
)
|
||
|
||
router = Router()
|
||
|
||
|
||
async def get_document_info(document_id: str, telegram_id: str):
|
||
"""Получить информацию о документе"""
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.get(
|
||
f"{settings.BACKEND_URL}/documents/{document_id}",
|
||
headers={"X-Telegram-ID": telegram_id}
|
||
) as response:
|
||
if response.status == 200:
|
||
return await response.json()
|
||
return None
|
||
except Exception as e:
|
||
print(f"Error getting document info: {e}")
|
||
return None
|
||
|
||
|
||
async def delete_document(document_id: str, telegram_id: str):
|
||
"""Удалить документ"""
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.delete(
|
||
f"{settings.BACKEND_URL}/documents/{document_id}",
|
||
headers={"X-Telegram-ID": telegram_id}
|
||
) as response:
|
||
return response.status == 204
|
||
except Exception as e:
|
||
print(f"Error deleting document: {e}")
|
||
return False
|
||
|
||
|
||
async def update_document(document_id: str, telegram_id: str, title: str = None, content: str = None):
|
||
"""Обновить документ"""
|
||
try:
|
||
update_data = {}
|
||
if title:
|
||
update_data["title"] = title
|
||
if content:
|
||
update_data["content"] = content
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.put(
|
||
f"{settings.BACKEND_URL}/documents/{document_id}",
|
||
json=update_data,
|
||
headers={"X-Telegram-ID": telegram_id}
|
||
) as response:
|
||
if response.status == 200:
|
||
return await response.json()
|
||
return None
|
||
except Exception as e:
|
||
print(f"Error updating document: {e}")
|
||
return None
|
||
|
||
|
||
async def upload_document_to_collection(collection_id: str, file_data: bytes, filename: str, telegram_id: str):
|
||
"""Загрузить документ в коллекцию"""
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
form_data = aiohttp.FormData()
|
||
form_data.add_field('file', file_data, filename=filename, content_type='application/octet-stream')
|
||
|
||
async with session.post(
|
||
f"{settings.BACKEND_URL}/documents/upload?collection_id={collection_id}",
|
||
data=form_data,
|
||
headers={"X-Telegram-ID": telegram_id}
|
||
) as response:
|
||
if response.status == 201:
|
||
return await response.json()
|
||
else:
|
||
error_text = await response.text()
|
||
print(f"Upload error: {response.status} - {error_text}")
|
||
return None
|
||
except Exception as e:
|
||
print(f"Error uploading document: {e}")
|
||
return None
|
||
|
||
|
||
@router.callback_query(lambda c: c.data.startswith("document:view:"))
|
||
async def view_document(callback: CallbackQuery):
|
||
"""Просмотр документа"""
|
||
document_id = callback.data.split(":")[2]
|
||
telegram_id = str(callback.from_user.id)
|
||
|
||
await callback.answer("Загружаю документ...")
|
||
|
||
document = await get_document_info(document_id, telegram_id)
|
||
if not document:
|
||
await callback.message.answer(
|
||
"<b>Ошибка</b>\n\nНе удалось загрузить документ.",
|
||
parse_mode="HTML"
|
||
)
|
||
return
|
||
|
||
title = document.get("title", "Без названия")
|
||
content = document.get("content", "")
|
||
collection_id = document.get("collection_id")
|
||
|
||
content_preview = content[:2000] if len(content) > 2000 else content
|
||
has_more = len(content) > 2000
|
||
|
||
response = f"<b>{title}</b>\n\n"
|
||
response += f"<i>{content_preview}</i>"
|
||
if has_more:
|
||
response += "\n\n<i>...</i>"
|
||
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.get(
|
||
f"{settings.BACKEND_URL}/collections/{collection_id}",
|
||
headers={"X-Telegram-ID": telegram_id}
|
||
) as response_collection:
|
||
if response_collection.status == 200:
|
||
collection_info = await response_collection.json()
|
||
owner_id = collection_info.get("owner_id")
|
||
|
||
async with session.get(
|
||
f"{settings.BACKEND_URL}/users/telegram/{telegram_id}"
|
||
) as response_user:
|
||
if response_user.status == 200:
|
||
user_info = await response_user.json()
|
||
current_user_id = user_info.get("user_id")
|
||
is_owner = str(owner_id) == str(current_user_id)
|
||
|
||
keyboard_buttons = []
|
||
if is_owner:
|
||
keyboard_buttons = [
|
||
[InlineKeyboardButton(text="Редактировать", callback_data=f"document:edit:{document_id}")],
|
||
[InlineKeyboardButton(text="Удалить", callback_data=f"document:delete:{document_id}")],
|
||
[InlineKeyboardButton(text="Назад", callback_data=f"collection:documents:{collection_id}")]
|
||
]
|
||
else:
|
||
keyboard_buttons = [
|
||
[InlineKeyboardButton(text="Редактировать", callback_data=f"document:edit:{document_id}")],
|
||
[InlineKeyboardButton(text="Назад", callback_data=f"collection:documents:{collection_id}")]
|
||
]
|
||
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=keyboard_buttons)
|
||
await callback.message.answer(response, parse_mode="HTML", reply_markup=keyboard)
|
||
return
|
||
except:
|
||
pass
|
||
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=[[
|
||
InlineKeyboardButton(text="Назад", callback_data=f"collection:documents:{collection_id}")
|
||
]])
|
||
await callback.message.answer(response, parse_mode="HTML", reply_markup=keyboard)
|
||
|
||
|
||
@router.callback_query(lambda c: c.data.startswith("document:edit:"))
|
||
async def edit_document_prompt(callback: CallbackQuery, state: FSMContext):
|
||
"""Запросить данные для редактирования документа"""
|
||
document_id = callback.data.split(":")[2]
|
||
telegram_id = str(callback.from_user.id)
|
||
|
||
document = await get_document_info(document_id, telegram_id)
|
||
if not document:
|
||
await callback.message.answer(
|
||
"<b>Ошибка</b>\n\nНе удалось загрузить документ.",
|
||
parse_mode="HTML"
|
||
)
|
||
await callback.answer()
|
||
return
|
||
|
||
await state.update_data(document_id=document_id)
|
||
await state.set_state(DocumentEditStates.waiting_for_title)
|
||
|
||
await callback.message.answer(
|
||
"<b>Редактирование документа</b>\n\n"
|
||
"Отправьте новое название документа или /skip чтобы оставить текущее.\n\n"
|
||
f"Текущее название: <b>{document.get('title', 'Без названия')}</b>",
|
||
parse_mode="HTML"
|
||
)
|
||
await callback.answer()
|
||
|
||
|
||
@router.message(StateFilter(DocumentEditStates.waiting_for_title))
|
||
async def process_edit_title(message: Message, state: FSMContext):
|
||
"""Обработать новое название документа"""
|
||
telegram_id = str(message.from_user.id)
|
||
data = await state.get_data()
|
||
document_id = data.get("document_id")
|
||
|
||
if message.text and message.text.strip() == "/skip":
|
||
new_title = None
|
||
else:
|
||
new_title = message.text.strip() if message.text else None
|
||
|
||
await state.update_data(title=new_title)
|
||
await state.set_state(DocumentEditStates.waiting_for_content)
|
||
|
||
await message.answer(
|
||
"<b>Содержимое документа</b>\n\n"
|
||
"Отправьте новое содержимое документа или /skip чтобы оставить текущее.",
|
||
parse_mode="HTML"
|
||
)
|
||
|
||
|
||
@router.message(StateFilter(DocumentEditStates.waiting_for_content))
|
||
async def process_edit_content(message: Message, state: FSMContext):
|
||
"""Обработать новое содержимое документа"""
|
||
telegram_id = str(message.from_user.id)
|
||
data = await state.get_data()
|
||
document_id = data.get("document_id")
|
||
title = data.get("title")
|
||
|
||
if message.text and message.text.strip() == "/skip":
|
||
new_content = None
|
||
else:
|
||
new_content = message.text.strip() if message.text else None
|
||
|
||
result = await update_document(document_id, telegram_id, title=title, content=new_content)
|
||
|
||
if result:
|
||
await message.answer(
|
||
"<b>Документ обновлен</b>\n\n"
|
||
"Изменения сохранены.",
|
||
parse_mode="HTML"
|
||
)
|
||
else:
|
||
await message.answer(
|
||
"<b>Ошибка</b>\n\n"
|
||
"Не удалось обновить документ.",
|
||
parse_mode="HTML"
|
||
)
|
||
|
||
await state.clear()
|
||
|
||
|
||
@router.callback_query(lambda c: c.data.startswith("document:delete:"))
|
||
async def delete_document_confirm(callback: CallbackQuery):
|
||
"""Подтверждение удаления документа"""
|
||
document_id = callback.data.split(":")[2]
|
||
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=[
|
||
[InlineKeyboardButton(text="Да, удалить", callback_data=f"document:delete_confirm:{document_id}")],
|
||
[InlineKeyboardButton(text="Отмена", callback_data=f"document:view:{document_id}")]
|
||
])
|
||
|
||
await callback.message.answer(
|
||
"<b>Подтверждение удаления</b>\n\n"
|
||
"Вы уверены, что хотите удалить этот документ?",
|
||
parse_mode="HTML",
|
||
reply_markup=keyboard
|
||
)
|
||
await callback.answer()
|
||
|
||
|
||
@router.callback_query(lambda c: c.data.startswith("document:delete_confirm:"))
|
||
async def delete_document_execute(callback: CallbackQuery):
|
||
"""Выполнить удаление документа"""
|
||
document_id = callback.data.split(":")[2]
|
||
telegram_id = str(callback.from_user.id)
|
||
|
||
await callback.answer("Удаляю документ...")
|
||
|
||
# Получаем информацию о документе для возврата к коллекции
|
||
document = await get_document_info(document_id, telegram_id)
|
||
collection_id = document.get("collection_id") if document else None
|
||
|
||
result = await delete_document(document_id, telegram_id)
|
||
|
||
if result:
|
||
await callback.message.answer(
|
||
"<b>Документ удален</b>",
|
||
parse_mode="HTML"
|
||
)
|
||
else:
|
||
await callback.message.answer(
|
||
"<b>Ошибка</b>\n\n"
|
||
"Не удалось удалить документ.",
|
||
parse_mode="HTML"
|
||
)
|
||
|
||
|
||
@router.callback_query(lambda c: c.data.startswith("document:upload:"))
|
||
async def upload_document_prompt(callback: CallbackQuery, state: FSMContext):
|
||
"""Запросить файл для загрузки"""
|
||
collection_id = callback.data.split(":")[2]
|
||
telegram_id = str(callback.from_user.id)
|
||
|
||
await state.update_data(collection_id=collection_id)
|
||
await state.set_state(DocumentUploadStates.waiting_for_file)
|
||
|
||
await callback.message.answer(
|
||
"<b>Загрузка документа</b>\n\n"
|
||
"Отправьте файл (PDF, PNG, JPG, JPEG, TIFF, BMP).\n\n"
|
||
"Поддерживаемые форматы:\n"
|
||
"• PDF\n"
|
||
"• Изображения: PNG, JPG, JPEG, TIFF, BMP",
|
||
parse_mode="HTML"
|
||
)
|
||
await callback.answer()
|
||
|
||
|
||
@router.message(StateFilter(DocumentUploadStates.waiting_for_file), F.document | F.photo)
|
||
async def process_upload_document(message: Message, state: FSMContext):
|
||
"""Обработать загрузку документа"""
|
||
telegram_id = str(message.from_user.id)
|
||
data = await state.get_data()
|
||
collection_id = data.get("collection_id")
|
||
|
||
if not collection_id:
|
||
await message.answer("Ошибка: не указана коллекция")
|
||
await state.clear()
|
||
return
|
||
|
||
file_id = None
|
||
filename = None
|
||
|
||
if message.document:
|
||
file_id = message.document.file_id
|
||
filename = message.document.file_name or "document.pdf"
|
||
|
||
supported_extensions = ['.pdf', '.png', '.jpg', '.jpeg', '.tiff', '.bmp']
|
||
file_ext = filename.lower().rsplit('.', 1)[-1] if '.' in filename else ''
|
||
if f'.{file_ext}' not in supported_extensions:
|
||
await message.answer(
|
||
"<b>Ошибка</b>\n\n"
|
||
f"Неподдерживаемый формат файла: {file_ext}\n\n"
|
||
"Поддерживаются: PDF, PNG, JPG, JPEG, TIFF, BMP",
|
||
parse_mode="HTML"
|
||
)
|
||
await state.clear()
|
||
return
|
||
elif message.photo:
|
||
file_id = message.photo[-1].file_id
|
||
filename = "photo.jpg"
|
||
else:
|
||
await message.answer(
|
||
"<b>Ошибка</b>\n\n"
|
||
"Пожалуйста, отправьте файл (PDF или изображение).",
|
||
parse_mode="HTML"
|
||
)
|
||
await state.clear()
|
||
return
|
||
|
||
try:
|
||
file = await message.bot.get_file(file_id)
|
||
file_data = await message.bot.download_file(file.file_path)
|
||
file_bytes = file_data.read()
|
||
|
||
result = await upload_document_to_collection(collection_id, file_bytes, filename, telegram_id)
|
||
|
||
if result:
|
||
await message.answer(
|
||
f"<b>Документ загружен</b>\n\n"
|
||
f"Название: <b>{result.get('title', filename)}</b>",
|
||
parse_mode="HTML"
|
||
)
|
||
else:
|
||
await message.answer(
|
||
"<b>Ошибка</b>\n\n"
|
||
"Не удалось загрузить документ.",
|
||
parse_mode="HTML"
|
||
)
|
||
except Exception as e:
|
||
print(f"Error uploading document: {e}")
|
||
await message.answer(
|
||
"<b>Ошибка</b>\n\n"
|
||
f"Произошла ошибка при загрузке: {str(e)}",
|
||
parse_mode="HTML"
|
||
)
|
||
|
||
await state.clear()
|
||
|