memes_bot/bot.py

import asyncio
import logging
import os
import uuid

from aiogram import Bot, Dispatcher, F
from aiogram.filters import Command, StateFilter
from aiogram.fsm.context import FSMContext
from aiogram.fsm.state import State, StatesGroup
from aiogram.fsm.storage.memory import MemoryStorage
from aiogram.types import (
    BotCommand,
    BotCommandScopeChat,
    CallbackQuery,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
    InlineQuery,
    InlineQueryResultCachedAudio,
    InlineQueryResultCachedDocument,
    InlineQueryResultCachedMpeg4Gif,
    InlineQueryResultCachedPhoto,
    InlineQueryResultCachedSticker,
    InlineQueryResultCachedVideo,
    InlineQueryResultCachedVoice,
    InlineQueryResultArticle,
    InputTextMessageContent,
    Message,
)

import db

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

BOT_TOKEN = os.environ["BOT_TOKEN"]

ALLOWED_USERS: set[int] = {
    int(uid.strip()) for uid in os.environ.get("ALLOWED_USERS", "").split(",") if uid.strip()
}
ADMIN_USERS: set[int] = {
    int(uid.strip()) for uid in os.environ.get("ADMIN_USERS", "").split(",") if uid.strip()
}

bot = Bot(token=BOT_TOKEN)
dp = Dispatcher(storage=MemoryStorage())

USER_COMMANDS = [
    BotCommand(command="start", description="Start the bot"),
    BotCommand(command="menu", description="Show available commands"),
    BotCommand(command="list", description="Browse saved memes"),
    BotCommand(command="cancel", description="Cancel current operation"),
]
ADMIN_COMMANDS = [
    *USER_COMMANDS,
    BotCommand(command="edit", description="Edit meme keywords"),
    BotCommand(command="delete", description="Remove a meme by ID"),
]


def is_allowed(user_id: int) -> bool:
    if user_id in ADMIN_USERS:
        return True
    return not ALLOWED_USERS or user_id in ALLOWED_USERS


def is_admin(user_id: int) -> bool:
    # If no ADMIN_USERS set, fall back to ALLOWED_USERS behaviour (backward compat)
    if ADMIN_USERS:
        return user_id in ADMIN_USERS
    return not ALLOWED_USERS or user_id in ALLOWED_USERS


async def register_commands() -> None:
    """Publish Telegram's command menu with role-appropriate suggestions."""
    if not ADMIN_USERS:
        await bot.set_my_commands(ADMIN_COMMANDS)
        return

    await bot.set_my_commands(USER_COMMANDS)
    for admin_id in ADMIN_USERS:
        await bot.set_my_commands(
            ADMIN_COMMANDS,
            scope=BotCommandScopeChat(chat_id=admin_id),
        )


class AddMedia(StatesGroup):
    waiting_for_keywords = State()


class EditMedia(StatesGroup):
    waiting_for_new_keywords = State()


def menu_text(user_id: int) -> str:
    if is_admin(user_id):
        return (
            "Send me any <b>photo, GIF, video, voice, music, sticker, or file</b>.\n"
            "I'll ask you for keywords next.\n\n"
            "/list — show all saved memes\n"
            "/edit [search] — edit keywords of a meme\n"
            "/delete &lt;id&gt; — remove a meme by ID\n"
            "/cancel — cancel current operation"
        )
    return (
        "Use me inline: type <code>@botname keyword</code> in any chat.\n\n"
        "/list — browse all saved memes"
    )


# ── /start ───────────────────────────────────────────────────────────────────

@dp.message(Command("start"))
async def cmd_start(msg: Message):
    if not is_allowed(msg.from_user.id):
        return
    await msg.answer(
        f"{menu_text(msg.from_user.id)}\n/menu — show this list again",
        parse_mode="HTML",
    )


# ── /menu ────────────────────────────────────────────────────────────────────

@dp.message(Command("menu"))
async def cmd_menu(msg: Message):
    if not is_allowed(msg.from_user.id):
        return
    await msg.answer(menu_text(msg.from_user.id), parse_mode="HTML")


# ── /cancel ──────────────────────────────────────────────────────────────────

@dp.message(Command("cancel"))
async def cmd_cancel(msg: Message, state: FSMContext):
    if not is_allowed(msg.from_user.id):
        return
    await state.clear()
    await msg.answer("Cancelled.")


# ── step 1: receive media (admin only) ───────────────────────────────────────

async def _receive_media(msg: Message, state: FSMContext, file_id: str, file_unique_id: str, media_type: str):
    await state.set_state(AddMedia.waiting_for_keywords)
    await state.update_data(
        file_id=file_id,
        file_unique_id=file_unique_id,
        media_type=media_type,
        keywords=[],
    )
    await msg.reply(
        "Got it! Now send keywords — one per message or comma-separated.\n"
        "Send /done when finished.",
        parse_mode="HTML",
    )


# Ordered: Telegram sets `document` alongside `animation` for backward
# compatibility, so animations must be matched before documents or every GIF
# would be filed as a file. Round videos are stored as plain "video" so they
# stay reachable from inline search, which has no cached video_note result.
MEDIA_EXTRACTORS = (
    ("photo",    lambda m: m.photo[-1] if m.photo else None),
    ("gif",      lambda m: m.animation),
    ("sticker",  lambda m: m.sticker),
    ("video",    lambda m: m.video),
    ("video",    lambda m: m.video_note),
    ("voice",    lambda m: m.voice),
    ("audio",    lambda m: m.audio),
    ("document", lambda m: m.document),
)

MEDIA_FILTER = (
    F.photo | F.animation | F.sticker | F.video
    | F.video_note | F.voice | F.audio | F.document
)


def extract_media(msg: Message) -> tuple[str, str, str] | None:
    """Return (media_type, file_id, file_unique_id) for any supported media."""
    for media_type, getter in MEDIA_EXTRACTORS:
        obj = getter(msg)
        if obj is not None:
            return media_type, obj.file_id, obj.file_unique_id
    return None


@dp.message(MEDIA_FILTER, StateFilter(None))
async def on_media(msg: Message, state: FSMContext):
    if not is_admin(msg.from_user.id):
        return
    found = extract_media(msg)
    if found is None:
        return
    media_type, file_id, file_unique_id = found
    await _receive_media(msg, state, file_id, file_unique_id, media_type)


# ── step 2: collect keywords ─────────────────────────────────────────────────

@dp.message(AddMedia.waiting_for_keywords, Command("done"))
async def cmd_done(msg: Message, state: FSMContext):
    data = await state.get_data()
    keywords: list[str] = data.get("keywords", [])
    if not keywords:
        await msg.answer("You haven't added any keywords yet. Send some, then /done.")
        return
    await db.add_media(
        msg.from_user.id,
        data["file_id"],
        data["file_unique_id"],
        data["media_type"],
        keywords,
    )
    await state.clear()
    await msg.answer(f"Saved with keywords: <b>{', '.join(keywords)}</b>", parse_mode="HTML")


@dp.message(AddMedia.waiting_for_keywords, F.text)
async def on_keyword(msg: Message, state: FSMContext):
    new_kws = [k.strip().lower() for k in msg.text.split(",") if k.strip()]
    if not new_kws:
        return
    data = await state.get_data()
    existing: list[str] = data.get("keywords", [])
    merged = existing + [k for k in new_kws if k not in existing]
    await state.update_data(keywords=merged)
    await msg.reply(f"Keywords so far: <b>{', '.join(merged)}</b>\nSend more or /done.", parse_mode="HTML")


@dp.message(AddMedia.waiting_for_keywords, MEDIA_FILTER)
async def on_media_during_keywords(msg: Message):
    await msg.reply("Still waiting for keywords. Send /done to save the current item, or /cancel to discard.")


# ── /list ────────────────────────────────────────────────────────────────────

LIST_PAGE_SIZE = 30


async def list_page(page: int) -> tuple[str, InlineKeyboardMarkup | None]:
    total = await db.count_media()
    if total == 0:
        return "No memes saved yet.", None

    total_pages = (total + LIST_PAGE_SIZE - 1) // LIST_PAGE_SIZE
    page = max(0, min(page, total_pages - 1))
    items = await db.list_media(page, LIST_PAGE_SIZE)
    lines = [f"<code>{item['id']}</code> [{item['media_type']}] — {item['keywords']}" for item in items]
    text = f"Page {page + 1} of {total_pages}\n" + "\n".join(lines)

    buttons = []
    if page > 0:
        buttons.append(InlineKeyboardButton(text="‹ Previous", callback_data=f"list:{page - 1}"))
    if page < total_pages - 1:
        buttons.append(InlineKeyboardButton(text="Next ›", callback_data=f"list:{page + 1}"))
    keyboard = InlineKeyboardMarkup(inline_keyboard=[buttons]) if buttons else None
    return text, keyboard


@dp.message(Command("list"))
async def cmd_list(msg: Message):
    if not is_allowed(msg.from_user.id):
        return
    text, keyboard = await list_page(0)
    await msg.answer(text, parse_mode="HTML", reply_markup=keyboard)


@dp.callback_query(F.data.startswith("list:"))
async def on_list_page(callback: CallbackQuery):
    if not is_allowed(callback.from_user.id):
        await callback.answer("Not allowed.", show_alert=True)
        return

    try:
        page = int(callback.data.split(":", 1)[1])
    except ValueError:
        await callback.answer("Invalid page.", show_alert=True)
        return

    text, keyboard = await list_page(page)
    await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
    await callback.answer()


# ── /delete ──────────────────────────────────────────────────────────────────

@dp.message(Command("delete"))
async def cmd_delete(msg: Message):
    if not is_admin(msg.from_user.id):
        await msg.answer("Only admins can delete memes.")
        return
    parts = msg.text.split(maxsplit=1)
    if len(parts) < 2:
        await msg.answer("Usage: /delete &lt;id&gt; — use the numeric ID from /list", parse_mode="HTML")
        return
    try:
        media_id = int(parts[1].strip())
    except ValueError:
        await msg.answer("Invalid ID. Use the numeric ID from /list.")
        return
    removed = await db.delete_media(media_id)
    await msg.answer("Deleted." if removed else "Not found.")


# ── /edit ────────────────────────────────────────────────────────────────────

@dp.message(Command("edit"))
async def cmd_edit(msg: Message):
    if not is_admin(msg.from_user.id):
        await msg.answer("Only admins can edit memes.")
        return
    parts = msg.text.split(maxsplit=1)
    query = parts[1].strip() if len(parts) > 1 else ""
    items = await db.search_media(query)
    if not items:
        no_match = f" matching '{query}'" if query else ""
        await msg.answer(f"No memes found{no_match}.")
        return
    buttons = []
    for item in items[:10]:
        label = f"[{item['media_type']}] {item['keywords'][:45]}"
        buttons.append([InlineKeyboardButton(text=label, callback_data=f"edit:{item['id']}")])
    keyboard = InlineKeyboardMarkup(inline_keyboard=buttons)
    await msg.answer("Select a meme to edit its keywords:", reply_markup=keyboard)


@dp.callback_query(F.data.startswith("edit:"))
async def on_edit_select(callback: CallbackQuery, state: FSMContext):
    if not is_admin(callback.from_user.id):
        await callback.answer("Not allowed.", show_alert=True)
        return
    media_id = int(callback.data.split(":", 1)[1])
    item = await db.get_media_by_id(media_id)
    if not item:
        await callback.answer("Meme not found.", show_alert=True)
        return
    await state.set_state(EditMedia.waiting_for_new_keywords)
    await state.update_data(media_id=media_id, keywords=[])
    await callback.message.edit_reply_markup(reply_markup=None)
    await callback.message.answer(
        f"Editing <b>#{item['id']}</b> [{item['media_type']}]\n"
        f"Current keywords: <b>{item['keywords']}</b>\n\n"
        "Send the new keywords (comma-separated or one per message).\n"
        "Send /done when finished, or /cancel to abort.",
        parse_mode="HTML",
    )
    await callback.answer()


@dp.message(EditMedia.waiting_for_new_keywords, Command("done"))
async def cmd_edit_done(msg: Message, state: FSMContext):
    data = await state.get_data()
    keywords: list[str] = data.get("keywords", [])
    if not keywords:
        await msg.answer("You haven't entered any keywords yet. Send some, then /done.")
        return
    updated = await db.update_media_keywords(data["media_id"], keywords)
    await state.clear()
    if updated:
        await msg.answer(f"Updated keywords: <b>{', '.join(keywords)}</b>", parse_mode="HTML")
    else:
        await msg.answer("Failed to update (meme may have been deleted).")


@dp.message(EditMedia.waiting_for_new_keywords, Command("cancel"))
async def cmd_edit_cancel(msg: Message, state: FSMContext):
    await state.clear()
    await msg.answer("Edit cancelled.")


@dp.message(EditMedia.waiting_for_new_keywords, F.text)
async def on_edit_keyword(msg: Message, state: FSMContext):
    new_kws = [k.strip().lower() for k in msg.text.split(",") if k.strip()]
    if not new_kws:
        return
    data = await state.get_data()
    existing: list[str] = data.get("keywords", [])
    merged = existing + [k for k in new_kws if k not in existing]
    await state.update_data(keywords=merged)
    await msg.reply(f"Keywords so far: <b>{', '.join(merged)}</b>\nSend more or /done.", parse_mode="HTML")


# Cached inline results, keyed by the media_type stored in the database.
# Audio carries its own title metadata; video and document require one.
INLINE_BUILDERS = {
    "photo": lambda rid, fid, kw: InlineQueryResultCachedPhoto(
        id=rid, photo_file_id=fid, title=kw, description=kw),
    "gif": lambda rid, fid, kw: InlineQueryResultCachedMpeg4Gif(
        id=rid, mpeg4_file_id=fid, title=kw),
    "video": lambda rid, fid, kw: InlineQueryResultCachedVideo(
        id=rid, video_file_id=fid, title=kw, description=kw),
    "voice": lambda rid, fid, kw: InlineQueryResultCachedVoice(
        id=rid, voice_file_id=fid, title=kw),
    "sticker": lambda rid, fid, kw: InlineQueryResultCachedSticker(
        id=rid, sticker_file_id=fid),
    "audio": lambda rid, fid, kw: InlineQueryResultCachedAudio(
        id=rid, audio_file_id=fid),
    "document": lambda rid, fid, kw: InlineQueryResultCachedDocument(
        id=rid, document_file_id=fid, title=kw or "file", description=kw),
}


# ── inline mode ──────────────────────────────────────────────────────────────

@dp.inline_query()
async def inline_handler(query: InlineQuery):
    if not is_allowed(query.from_user.id):
        await query.answer([], cache_time=1)
        return

    try:
        results_raw = await db.search_media(query.query)
        results = []

        for item in results_raw:
            rid = str(uuid.uuid4())
            kw_display = item["keywords"].replace(",", ", ")
            mtype = item["media_type"]

            builder = INLINE_BUILDERS.get(mtype)
            if builder is not None:
                results.append(builder(rid, item["file_id"], kw_display))

        if not results:
            results = [InlineQueryResultArticle(
                id="noop", title="No results found",
                input_message_content=InputTextMessageContent(message_text="—"),
                description=f"No memes matching '{query.query}'",
            )]

        await query.answer(results, cache_time=5, is_personal=True)

    except Exception:
        log.exception("inline_handler error for user %s query %r", query.from_user.id, query.query)
        await query.answer([], cache_time=1)


# ── main ─────────────────────────────────────────────────────────────────────

async def main():
    await db.init_db()
    await register_commands()
    log.info("Bot starting…")
    await dp.start_polling(bot)


if __name__ == "__main__":
    asyncio.run(main())