mfinder/plugins/index.py

import asyncio
from pyrogram import Client, filters
from pyrogram.errors import FloodWait
from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup
from mfinder import ADMINS, LOGGER
from mfinder.db.files_sql import save_file, delete_file
from mfinder.utils.helpers import edit_caption

lock = asyncio.Lock()
media_filter = filters.document | filters.video | filters.audio
DELETING_USERS = set()

@Client.on_message(filters.private & filters.user(ADMINS) & media_filter)
async def index_files(bot, message):
    user_id = message.from_user.id
    if user_id in DELETING_USERS:
        try:
            for file_type in ("document", "video", "audio"):
                media = getattr(message, file_type, None)
                if not media:
                    continue
                del_file = await delete_file(media)
                if del_file == "Not Found":
                    await message.reply(f"❌ `{media.file_name}` not found in database.")
                elif del_file == True:
                    await message.reply(f"🗑️ `{media.file_name}` deleted from database successfully!")
                else:
                    await message.reply(f"⚠️ Error occurred while deleting `{media.file_name}`.")
        except Exception as e:
            LOGGER.warning("Error occurred during bulk delete: %s", str(e))
        return

    if lock.locked():
        await message.reply("Wait until previous process complete.")
    else:
        try:
            last_msg_id = None
            chat_id = None
            if hasattr(message, "forward_origin") and message.forward_origin:
                origin = message.forward_origin
                if hasattr(origin, "chat") and origin.chat:
                    chat_id = origin.chat.username or origin.chat.id
                elif hasattr(origin, "sender_chat") and origin.sender_chat:
                    chat_id = origin.sender_chat.username or origin.sender_chat.id
                elif hasattr(origin, "sender_user") and origin.sender_user:
                    chat_id = origin.sender_user.id
                if hasattr(origin, "message_id"):
                    last_msg_id = origin.message_id
            if not last_msg_id:
                last_msg_id = getattr(message, "forward_from_message_id", None)
            if not chat_id:
                from_chat = getattr(message, "forward_from_chat", None)
                if from_chat:
                    chat_id = from_chat.username or from_chat.id
            if not last_msg_id or not chat_id:
                raise ValueError("Could not extract chat details from the forwarded message.")
            await bot.get_messages(chat_id, last_msg_id)

            kb = InlineKeyboardMarkup(
                [
                    [InlineKeyboardButton("Proceed", callback_data=f"index {chat_id} {last_msg_id}")],
                    [InlineKeyboardButton("Cancel", callback_data="can-index")],
                ]
            )
            await bot.send_message(user_id, "Please confirm if you want to start indexing", reply_markup=kb)
        except Exception as e:
            await message.reply_text(f"Unable to start indexing. Error: <code>{e}</code>")

@Client.on_callback_query(filters.regex(r"^index -?\d+ \d+"))
async def index(bot, query):
    user_id = query.from_user.id
    chat_id, last_msg_id = map(int, query.data.split()[1:])
    await query.message.delete()
    msg = await bot.send_message(user_id, "Processing Index...⏳")
    total_files = 0
    async with lock:
        try:
            total = last_msg_id + 1
            current = 2
            counter = 0
            while True:
                try:
                    message = await bot.get_messages(chat_id=chat_id, message_ids=current, replies=0)
                except FloodWait as e:
                    LOGGER.warning("FloodWait while indexing, Error: %s", str(e))
                    await asyncio.sleep(e.value)
                except Exception as e:
                    LOGGER.warning("Error occurred while fetching message: %s", str(e))
                try:
                    for file_type in ("document", "video", "audio"):
                        media = getattr(message, file_type, None)
                        if not media:
                            break
                        file_name = media.file_name
                        file_name = edit_caption(file_name)
                        media.file_type = file_type
                        media.caption = message.caption if message.caption else file_name
                        await save_file(media)
                        total_files += 1
                        try:
                            from mfinder.plugins.requests import check_and_notify_pending_requests
                            asyncio.create_task(check_and_notify_pending_requests(bot, file_name))
                        except Exception as req_err:
                            LOGGER.warning(f"Failed to launch check_and_notify_pending_requests in index: {req_err}")
                except Exception as e:
                    LOGGER.warning("Error occurred while saving file: %s", str(e))

                current += 1
                counter += 1
                if counter == 50:
                    try:
                        await msg.edit(f"Total messages fetched: {current}\nTotal messages saved: {total_files}")
                    except FloodWait as e:
                        LOGGER.warning("FloodWait while indexing, sleeping for: %s", str(e.value))
                        await asyncio.sleep(e.value)
                    counter -= 50
                if current == total:
                    break
        except Exception as e:
            LOGGER.exception(e)
            await msg.edit(f"Error: {e}")
        else:
            await msg.edit(f"Total {total_files} Saved To DataBase!")

@Client.on_message(filters.command(["index"]) & filters.user(ADMINS))
async def index_comm(bot, update):
    await update.reply("Now please forward the last message of the channel you want to index & follow the steps. Bot must be admin of the channel if the channel is private.")

@Client.on_message(filters.command(["delete"]) & filters.user(ADMINS))
async def delete_files(bot, message):
    if message.reply_to_message:
        org_msg = message.reply_to_message
        try:
            for file_type in ("document", "video", "audio"):
                media = getattr(org_msg, file_type, None)
                if not media:
                    break
                del_file = await delete_file(media)
                if del_file == "Not Found":
                    await message.reply(f"❌ `{media.file_name}` not found in database")
                elif del_file == True:
                    await message.reply(f"🗑️ `{media.file_name}` deleted from database")
                else:
                    await message.reply(f"⚠️ Error occurred while deleting `{media.file_name}`")
        except Exception as e:
            LOGGER.warning("Error occurred while deleting file: %s", str(e))
    else:
        user_id = message.from_user.id
        DELETING_USERS.add(user_id)
        await message.reply(
            "🧹 **Bulk Delete Mode Activated!**\n\n"
            "Please send or forward all the files you want to delete from the database. "
            "I will delete them one-by-one automatically.\n\n"
            "When you are finished, type `/done` to turn off this mode."
        )

@Client.on_message(filters.command(["done"]) & filters.user(ADMINS))
async def done_deleting(bot, message):
    user_id = message.from_user.id
    if user_id in DELETING_USERS:
        DELETING_USERS.remove(user_id)
        await message.reply("🧹 **Bulk Delete Mode Deactivated!**")
    else:
        await message.reply("❌ You are not currently in bulk delete mode.")

@Client.on_callback_query(filters.regex(r"^can-index$"))
async def cancel_index(bot, query):
    await query.message.delete()