qbittorent_bot/bot.py

import logging
import os
from datetime import datetime

# Импорты для Telegram Bot API
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (
    Application,
    CommandHandler,
    MessageHandler,
    filters,
    ContextTypes,
    CallbackQueryHandler,
)

# Импорты для qBittorrent API
from qbittorrentapi import Client, APIError

# --- Настройка логирования ---
logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
)
logger = logging.getLogger(__name__)

# --- Переменные окружения для qBittorrent ---
TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
QBT_HOST = os.getenv("QBT_HOST")
QBT_PORT = os.getenv("QBT_PORT")
QBT_USERNAME = os.getenv("QBT_USERNAME")
QBT_PASSWORD = os.getenv("QBT_PASSWORD")

# --- Глобальная переменная для qBittorrent клиента ---
qb = None

# --- Инициализация qBittorrent клиента ---
def init_qbittorrent_client():
    global qb
    if not all([QBT_HOST, QBT_PORT, QBT_USERNAME, QBT_PASSWORD]):
        logger.error("QBittorrent credentials (host, port, username, password) are not fully set in environment variables.")
        return False
    try:
        qb = Client(
            host=f"{QBT_HOST}:{QBT_PORT}",
            username=QBT_USERNAME,
            password=QBT_PASSWORD
        )
        qb_version = qb.app.version
        logger.info(f"Connected to qBittorrent API v{qb_version} on {QBT_HOST}:{QBT_PORT}")
        return True
    except APIError as e:
        logger.error(f"Failed to connect or login to qBittorrent: {e}. Check your qBittorrent Web UI address and credentials.")
        qb = None
        return False
    except Exception as e:
        logger.error(f"An unexpected error occurred during qBittorrent connection: {e}")
        qb = None
        return False

# --- Команда /start ---
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if update.message is None:
        logger.warning("Received an update without a message object in start handler.")
        return

    logger.info(f"Received /start command from {update.effective_user.id}")

    if not init_qbittorrent_client():
        await update.message.reply_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return

    await update.message.reply_text(
        "Привет! Я бот для управления qBittorrent.\n"
        "Отправь мне magnet-ссылку или URL torrent-файла, "
        "чтобы добавить загрузку.\n"
        "Используй /status для просмотра текущих загрузок.\n"
        "Используй /stop_torrent для остановки загрузки.\n"
        "Используй /help для списка команд."
    )

# --- Команда /help ---
async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if update.message is None:
        logger.warning("Received an update without a message object in help handler.")
        return

    logger.info(f"Received /help command from {update.effective_user.id}")

    help_text = (
        "Вот список доступных команд:\n\n"
        "/start - Начать работу с ботом и проверить подключение к qBittorrent.\n"
        "/status - Показать текущий статус всех активных загрузок и управлять ими.\n" 
        "/stop_torrent - Выбрать и остановить загрузку торрента (устаревает, используйте /status).\n"
        "/help - Показать это справочное сообщение.\n\n"
        "Также вы можете отправить мне magnet-ссылку или URL torrent-файла " 
        "для добавления загрузки. Бот предложит выбрать категорию и директорию."
    )
    await update.message.reply_text(help_text) 


# --- Команда /status ---
async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if update.message is None:
        logger.warning("Received an update without a message object in status handler.")
        return

    logger.info(f"Received /status command from {update.effective_user.id}")

    if not init_qbittorrent_client():
        await update.message.reply_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return

    try:
        # Получаем общую статистику qBittorrent
        global_transfer_info = qb.transfer_info()
        
        global_dlspeed_bytes = global_transfer_info.dl_info_speed
        global_upspeed_bytes = global_transfer_info.up_info_speed

        # Форматирование скорости
        def format_speed(speed_bytes):
            if speed_bytes >= (1024 * 1024):
                return f"{speed_bytes / (1024 * 1024):.2f} MB/s"
            elif speed_bytes >= 1024:
                return f"{speed_bytes / 1024:.2f} KB/s"
            else:
                return f"{speed_bytes} B/s"

        global_dlspeed_formatted = format_speed(global_dlspeed_bytes)
        global_upspeed_formatted = format_speed(global_upspeed_bytes)

        torrents = qb.torrents_info()
        
        # Считаем торренты по состояниям
        active_count = sum(1 for t in torrents if t.state in ['downloading', 'stalledDL', 'uploading', 'checkingQT', 'queuedDL', 'checkingUP', 'queuedUP'])
        paused_count = sum(1 for t in torrents if t.state in ['pausedDL', 'pausedUP', 'stoppedUP', 'stoppedDL'])
        # completed и seeding - состояния завершенных загрузок
        completed_count = sum(1 for t in torrents if t.state in ['completed', 'seeding', 'stalledUP']) 
        
        summary_text = (
            f"⚡️ *Общий статус qBittorrent*\n"
            f"   ⬇️ Общая скорость загрузки: {global_dlspeed_formatted}\n"
            f"   ⬆️ Общая скорость отдачи: {global_upspeed_formatted}\n"
            f"   Активных торрентов: {active_count}\n"
            f"   На паузе/Остановлено: {paused_count}\n"
            f"   Завершено/Раздается: {completed_count}\n"
            f"---\n"
        )
        await update.message.reply_text(summary_text, parse_mode="Markdown")


        if not torrents:
            await update.message.reply_text("Загрузок не найдено.")
            return

        for torrent in torrents:
            download_speed = torrent.dlspeed / (1024 * 1024) if torrent.dlspeed >= (1024 * 1024) else torrent.dlspeed / 1024
            upload_speed = torrent.upspeed / (1024 * 1024) if torrent.upspeed >= (1024 * 1024) else torrent.upspeed / 1024

            dl_unit = "MB/s" if torrent.dlspeed >= (1024 * 1024) else "KB/s"
            up_unit = "MB/s" if torrent.upspeed >= (1024 * 1024) else "KB/s"

            eta_str = ""
            if torrent.eta == 8640000:
                eta_str = "∞"
            elif torrent.eta > 0:
                hours, remainder = divmod(torrent.eta, 3600)
                minutes, seconds = divmod(remainder, 60)
                if hours > 0:
                    eta_str = f"{int(hours)}ч {int(minutes)}мин"
                elif minutes > 0:
                    eta_str = f"{int(minutes)}мин {int(seconds)}с"
                else:
                    eta_str = f"{int(seconds)}с"
            else:
                eta_str = "Завершено"


            message_text = (
                f"📊 *{torrent.name}*\n"
                f"   Состояние: {torrent.state}\n"
                f"   Прогресс: {torrent.progress:.2%}\n"
                f"   ⬇️ {download_speed:.2f} {dl_unit} ⬆️ {upload_speed:.2f} {up_unit}\n"
                f"   ETA: {eta_str}\n"
                f"   Размер: {(torrent.size / (1024*1024*1024)):.2f} ГБ"
            )
            
            control_buttons = []
            if torrent.state in ['downloading', 'stalledDL', 'uploading', 'checkingQT', 'queuedDL', 'checkingUP', 'queuedUP']:
                control_buttons.append(InlineKeyboardButton("🔴 Остановить", callback_data=f"stop_hash_{torrent.hash}"))
            elif torrent.state in ['pausedDL', 'pausedUP', 'stoppedUP', 'stoppedDL']: 
                control_buttons.append(InlineKeyboardButton("▶️ Запустить", callback_data=f"start_hash_{torrent.hash}"))
            elif torrent.state == 'metaDL': 
                 pass 
            else:
                control_buttons.append(InlineKeyboardButton("ℹ️ Неизвестное состояние", callback_data=f"info_hash_{torrent.hash}"))

            delete_buttons = [
                InlineKeyboardButton("🗑️ Удалить (без файлов)", callback_data=f"delete_hash_{torrent.hash}_false"),
                InlineKeyboardButton("❌ Удалить (с файлами)", callback_data=f"delete_hash_{torrent.hash}_true")
            ]

            keyboard_rows = []
            if control_buttons: # Добавляем кнопки управления, если они есть
                keyboard_rows.append(control_buttons)
            keyboard_rows.append(delete_buttons) # Всегда добавляем кнопки удаления

            reply_markup = InlineKeyboardMarkup(keyboard_rows) if keyboard_rows else None
            
            await update.message.reply_text(
                message_text,
                parse_mode="Markdown",
                reply_markup=reply_markup
            )

    except APIError as e:
        logger.error(f"Error getting torrent status: {e}")
        await update.message.reply_text(f"Ошибка при получении статуса торрентов: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred in status command: {e}")
        await update.message.reply_text(f"Произошла непредвиденная ошибка: {e}")


# --- Команда /stop_torrent (Теперь это будет устаревшая команда) ---
async def stop_torrent(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if update.message is None:
        logger.warning("Received an update without a message object in stop_torrent handler.")
        return

    logger.info(f"Received /stop_torrent command from {update.effective_user.id}")

    await update.message.reply_text(
        "Используйте команду /status для управления торрентами через кнопки 'Запустить' и 'Остановить'."
    )


# --- Обработка кнопки "Остановить торрент" ---
async def stop_torrent_callback(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.callback_query
    await query.answer()

    torrent_hash = query.data.replace("stop_hash_", "")
    logger.info(f"Attempting to pause torrent with hash: {torrent_hash}") 

    if not init_qbittorrent_client():
        await query.edit_message_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return

    try:
        qb.torrents_pause(torrent_hashes=torrent_hash) 
        await query.edit_message_text(f"Торрент ({torrent_hash[:6]}...) успешно остановлен (поставлен на паузу).")
    except APIError as e:
        logger.error(f"Error pausing torrent {torrent_hash}: {e}")
        await query.edit_message_text(f"Ошибка при остановке торрента: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred during torrent pausing: {e}")
        await query.edit_message_text(f"Произошла непредвиденная ошибка: {e}")

# --- Обработка кнопки "Запустить торрент" ---
async def start_torrent_callback(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.callback_query
    await query.answer()

    torrent_hash = query.data.replace("start_hash_", "")
    logger.info(f"Attempting to resume torrent with hash: {torrent_hash}") 

    if not init_qbittorrent_client():
        await query.edit_message_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return

    try:
        qb.torrents_resume(torrent_hashes=torrent_hash)
        await query.edit_message_text(f"Торрент ({torrent_hash[:6]}...) успешно запущен (возобновлен).")
    except APIError as e:
        logger.error(f"Error resuming torrent {torrent_hash}: {e}")
        await query.edit_message_text(f"Ошибка при запуске торрента: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred during torrent resuming: {e}")
        await query.edit_message_text(f"Произошла непредвиденная ошибка: {e}")

# --- Обработка кнопки "Удалить торрент" ---
async def delete_torrent_callback(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.callback_query
    await query.answer()

    parts = query.data.split('_')
    torrent_hash = parts[2]
    delete_files_str = parts[3]
    delete_files = True if delete_files_str == 'true' else False

    logger.info(f"Attempting to delete torrent with hash: {torrent_hash}, delete_files: {delete_files}")

    if not init_qbittorrent_client():
        await query.edit_message_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return

    try:
        qb.torrents_delete(torrent_hashes=torrent_hash, delete_files=delete_files)
        action_text = "с файлами" if delete_files else "без файлов"
        await query.edit_message_text(f"Торрент ({torrent_hash[:6]}...) успешно удален {action_text}.")
    except APIError as e:
        logger.error(f"Error deleting torrent {torrent_hash}: {e}")
        await query.edit_message_text(f"Ошибка при удалении торрента: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred during torrent deletion: {e}")
        await query.edit_message_text(f"Произошла непредвиденная ошибка при удалении торрента: {e}")


# --- Обработка magnet-ссылок и URL ---
async def handle_url(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if update.message is None:
        logger.warning("Received an update without a message object in handle_url handler.")
        return

    text = update.message.text
    logger.info(f"Received URL/Magnet: {text} from {update.effective_user.id}")

    if not init_qbittorrent_client():
        await update.message.reply_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return

    try:
        context.user_data['current_torrent_url'] = text

        categories_dict = qb.torrents_categories()
        category_keyboard = []
        for category_name in categories_dict.keys():
            category_keyboard.append([InlineKeyboardButton(category_name, callback_data=f"select_category_{category_name}")])
        
        category_keyboard.append([InlineKeyboardButton("Без категории", callback_data="select_category_no_category")])

        reply_markup = InlineKeyboardMarkup(category_keyboard)
        await update.message.reply_text('Выберите категорию для загрузки:', reply_markup=reply_markup)

    except APIError as e:
        logger.error(f"Error fetching categories: {e}")
        await update.message.reply_text(f"Ошибка при получении категорий: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred in handle_url: {e}")
        await update.message.reply_text(f"Произошла непредвиденная ошибка: {e}")


# --- Обработка выбора категории ---
async def select_category_callback(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.callback_query
    await query.answer()

    data = query.data.split('_')
    selected_category = data[2] if len(data) > 2 and data[2] != 'no_category' else None

    logger.info(f"Selected category: {selected_category}")

    context.user_data['selected_category'] = selected_category

    await query.edit_message_text("Выберите директорию загрузки:")
    await send_directory_options(query, context)

# --- Отправка опций директорий ---
async def send_directory_options(query, context):
    if not init_qbittorrent_client():
        await query.edit_message_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return

    try:
        available_paths = [
            qb.app.default_save_path,
            # --- ВАЖНО: ЗАМЕНИТЕ ЭТИ ПУТИ НА СВОИ АКТУАЛЬНЫЕ ПУТИ, ВИДИМЫЕ QBITTORRENT ---
            "/share/Data/Films", 
            "/share/Data/Serials", # Пример пути
            "/share/Data/torrents", # Пример пути
            #"/var/lib/qbittorrent/data/completed", # Пример пути
            # -------------------------------------------------------------------------
        ]
        
        # Удаляем дубликаты и сортируем пути для удобства
        available_paths = sorted(list(set([p.replace(os.sep, '/') for p in available_paths if p])))
        
        directory_keyboard = []
        for path in available_paths:
            # Обрезаем путь для отображения, если он слишком длинный
            display_path = os.path.basename(path) if len(path) > 30 else path
            directory_keyboard.append([InlineKeyboardButton(display_path, callback_data=f"select_dir_{path}")])
        
        reply_markup = InlineKeyboardMarkup(directory_keyboard)
        await query.edit_message_reply_markup(reply_markup=reply_markup)
    
    except APIError as e:
        logger.error(f"Error fetching default save path: {e}")
        await query.edit_message_text(f"Ошибка при получении директорий сохранения: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred in send_directory_options: {e}")
        await query.edit_message_text(f"Произошла непредвиденная ошибка: {e}")


# --- Обработка выбора директории и добавление торрента ---
async def select_directory_callback(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.callback_query
    await query.answer()

    data = query.data.split('_')
    # Собираем путь обратно из частей, так как он может содержать '_'
    selected_directory = "_".join(data[2:]) 

    logger.info(f"Selected directory: {selected_directory}")

    torrent_url = context.user_data.get('current_torrent_url')
    category = context.user_data.get('selected_category')

    if not torrent_url:
        await query.edit_message_text("Ошибка: URL торрента не найден. Пожалуйста, попробуйте снова.")
        return

    if not init_qbittorrent_client():
        await query.edit_message_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return

    try:
        # Логирование параметров перед добавлением торрента
        logger.info(f"Attempting to add torrent:")
        logger.info(f"  URL: {torrent_url}")
        logger.info(f"  Category: {category or 'None'}")
        logger.info(f"  Save Path: {selected_directory}")

        qb.torrents_add(
            urls=torrent_url,
            category=category,
            save_path=selected_directory
        )
        await query.edit_message_text(
            f"Торрент успешно добавлен в qBittorrent.\n"
            f"Категория: {category or 'Без категории'}\n"
            f"Директория: {selected_directory}"
        )
        # Очистка user_data после успешного добавления
        if 'current_torrent_url' in context.user_data:
            del context.user_data['current_torrent_url']
        if 'selected_category' in context.user_data:
            del context.user_data['selected_category']
    except APIError as e:
        logger.error(f"API Error adding torrent: {e}")
        await query.edit_message_text(f"Ошибка API при добавлении торрента: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred during torrent addition: {e}")
        await query.edit_message_text(f"Произошла непредвиденная ошибка при добавлении торрента: {e}")


# --- Обработчик ошибок Telegram Bot API ---
async def error_handler(update: object, context: ContextTypes.DEFAULT_TYPE) -> None:
    logger.error(f"Exception while handling an update: {context.error}")
    if update and update.effective_message:
        await update.effective_message.reply_text(f"Произошла внутренняя ошибка бота: {context.error}")

# --- Обработчик для неизвестных команд ---
async def unknown_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if update.message is None:
        logger.warning("Received an unknown command update without a message object.")
        return
    logger.info(f"Received unknown command: {update.message.text} from {update.effective_user.id}")
    await update.message.reply_text("Извините, я не понял эту команду.")

# --- Обработчик для любого другого текста ---
async def echo(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if update.message is None:
        logger.warning(f"Received non-text update in echo handler: {update}")
        return
    logger.info(f"Received non-command text: {update.message.text} from {update.effective_user.id}")
    await update.message.reply_text(f"Вы сказали: {update.message.text}")


# --- Основная функция, запускающая бота ---
def main() -> None:
    if not TELEGRAM_BOT_TOKEN:
        logger.critical("TELEGRAM_BOT_TOKEN environment variable is not set. Exiting.")
        exit(1)

    application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()

    # --- Добавление обработчиков команд ---
    application.add_handler(CommandHandler("start", start))
    application.add_handler(CommandHandler("status", status))
    application.add_handler(CommandHandler("stop_torrent", stop_torrent)) 
    application.add_handler(CommandHandler("help", help_command))

    # --- Добавление обработчиков сообщений ---
    url_regex = r"magnet:\?xt=urn:[a-z0-9]+"
    torrent_url_regex = r"https?://[^\s]+(?:/\d+|/\w+\.torrent|/download/\d+|\/torrent\.php\?hash=)[\S]*"
    
    application.add_handler(MessageHandler(filters.TEXT & (filters.Regex(url_regex) | filters.Regex(torrent_url_regex)), handle_url))
    
    # Добавляем CallbackQueryHandler для кнопок выбора категории
    application.add_handler(CallbackQueryHandler(select_category_callback, pattern=r"^select_category_.*"))
    # Добавляем CallbackQueryHandler для кнопок выбора директории
    application.add_handler(CallbackQueryHandler(select_directory_callback, pattern=r"^select_dir_.*"))
    # CallbackQueryHandler для кнопок остановки торрентов
    application.add_handler(CallbackQueryHandler(stop_torrent_callback, pattern=r"^stop_hash_.*"))
    # CallbackQueryHandler для кнопок запуска торрентов
    application.add_handler(CallbackQueryHandler(start_torrent_callback, pattern=r"^start_hash_.*"))
    # CallbackQueryHandler для кнопок удаления торрентов
    application.add_handler(CallbackQueryHandler(delete_torrent_callback, pattern=r"^delete_hash_.*"))


    # --- Обработчик для любого другого текста, не являющегося командой ---
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, echo))

    # --- Обработчик для неизвестных команд ---
    application.add_handler(MessageHandler(filters.COMMAND, unknown_command))

    # --- Добавление обработчика ошибок (ИСПРАВЛЕНО!) ---
    application.add_error_handler(error_handler)

    # --- Запуск бота ---
    logger.info("Bot started polling...")
    application.run_polling(allowed_updates=Update.ALL_TYPES)

if __name__ == "__main__":
    main()