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()
}

# Off by default: with ALLOWED_USERS unset, is_allowed() lets everyone in, so
# opening submissions implicitly would accept media from the whole world.
ACCEPT_SUBMISSIONS = os.environ.get("ACCEPT_SUBMISSIONS", "").strip().lower() in {"1", "true", "yes"}

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="review", description="Review pending submissions"),
    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"
            "/review — review pending submissions\n"
            "/edit [search] — edit keywords of a meme\n"
            "/delete &lt;id&gt; — remove a meme by ID\n"
            "/cancel — cancel current operation"
        )
    if ACCEPT_SUBMISSIONS:
        return (
            "Use me inline: type <code>@botname keyword</code> in any chat.\n\n"
            "Send me media to submit it — a moderator will review it before it goes live.\n\n"
            "/list — browse all saved memes"
        )
    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, status: str = "approved"):
    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=[],
        status=status,
    )
    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):
    user_id = msg.from_user.id
    if is_admin(user_id):
        status = "approved"
    elif ACCEPT_SUBMISSIONS and is_allowed(user_id):
        status = "pending"
    else:
        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, status)


# ── 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
    status = data.get("status", "approved")
    await db.add_media(
        msg.from_user.id,
        data["file_id"],
        data["file_unique_id"],
        data["media_type"],
        keywords,
        status=status,
    )
    await state.clear()
    if status == "pending":
        await msg.answer(
            f"Submitted with keywords: <b>{', '.join(keywords)}</b>\n"
            "A moderator will review it before it becomes searchable.",
            parse_mode="HTML",
        )
    else:
        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)
    from_review = data.get("from_review", False)
    await state.clear()

    if not updated:
        await msg.answer("Failed to update (meme may have been deleted).")
        return

    if from_review:
        # The item is still pending, so the queue hands the same one back with
        # its new keywords — ready to approve.
        await send_next_pending(msg.chat.id, "Keywords updated.")
        return

    await msg.answer(f"Updated keywords: <b>{', '.join(keywords)}</b>", parse_mode="HTML")


@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")


# ── /review ──────────────────────────────────────────────────────────────────

# Sending a stored item back needs one call per media type. Stickers cannot
# carry a caption, so the item and its review buttons are sent as two messages
# rather than special-casing captions per type.
MEDIA_SENDERS = {
    "photo": lambda chat_id, fid: bot.send_photo(chat_id, fid),
    "gif": lambda chat_id, fid: bot.send_animation(chat_id, fid),
    "video": lambda chat_id, fid: bot.send_video(chat_id, fid),
    "voice": lambda chat_id, fid: bot.send_voice(chat_id, fid),
    "audio": lambda chat_id, fid: bot.send_audio(chat_id, fid),
    "sticker": lambda chat_id, fid: bot.send_sticker(chat_id, fid),
    "document": lambda chat_id, fid: bot.send_document(chat_id, fid),
}


def review_keyboard(media_id: int) -> InlineKeyboardMarkup:
    return InlineKeyboardMarkup(inline_keyboard=[
        [
            InlineKeyboardButton(text="✅ Approve", callback_data=f"review:approve:{media_id}"),
            InlineKeyboardButton(text="❌ Reject", callback_data=f"review:reject:{media_id}"),
        ],
        [
            InlineKeyboardButton(text="✏️ Fix keywords", callback_data=f"review:edit:{media_id}"),
            InlineKeyboardButton(text="⏸ Stop", callback_data="review:stop:0"),
        ],
    ])


async def send_next_pending(chat_id: int, notice: str = "") -> None:
    """Show the next item awaiting review, or report the queue is empty.

    Nothing tracks where a moderator got to — each call asks the database for
    the oldest pending item. Stopping and resuming later needs no state, and
    two moderators working at once simply take items off the same queue.
    """
    item = await db.next_pending()
    prefix = f"{notice}\n\n" if notice else ""

    if item is None:
        await bot.send_message(chat_id, f"{prefix}Nothing left to review.")
        return

    remaining = await db.count_pending()
    send = MEDIA_SENDERS.get(item["media_type"])
    if send is not None:
        try:
            await send(chat_id, item["file_id"])
        except Exception:
            # A stale or unusable file_id must not wedge the queue — show the
            # details anyway so the moderator can still reject it.
            log.exception("could not send pending item %s for review", item["id"])
            await bot.send_message(chat_id, "⚠️ Could not load this media.")
    else:
        log.warning("no sender for media_type %r (item %s)", item["media_type"], item["id"])
        await bot.send_message(chat_id, f"⚠️ Unsupported media type: {item['media_type']}")

    await bot.send_message(
        chat_id,
        f"{prefix}Pending <b>#{item['id']}</b> [{item['media_type']}]\n"
        f"From user <code>{item['owner_id']}</code>\n"
        f"Keywords: <b>{item['keywords'].replace(',', ', ')}</b>\n\n"
        f"{remaining} awaiting review.",
        parse_mode="HTML",
        reply_markup=review_keyboard(item["id"]),
    )


@dp.message(Command("review"))
async def cmd_review(msg: Message):
    if not is_admin(msg.from_user.id):
        await msg.answer("Only moderators can review submissions.")
        return
    await send_next_pending(msg.chat.id)


@dp.callback_query(F.data.startswith("review:"))
async def on_review_action(callback: CallbackQuery, state: FSMContext):
    if not is_admin(callback.from_user.id):
        await callback.answer("Not allowed.", show_alert=True)
        return

    _, action, raw_id = callback.data.split(":", 2)
    await callback.message.edit_reply_markup(reply_markup=None)

    if action == "stop":
        await callback.message.answer("Review paused. Send /review to pick up where you left off.")
        await callback.answer()
        return

    media_id = int(raw_id)

    if action == "edit":
        item = await db.get_media_by_id(media_id)
        if not item:
            await send_next_pending(callback.message.chat.id, "That item is gone.")
            await callback.answer()
            return
        await state.set_state(EditMedia.waiting_for_new_keywords)
        await state.update_data(media_id=media_id, keywords=[], from_review=True)
        await callback.message.answer(
            f"Editing <b>#{item['id']}</b>\n"
            f"Current keywords: <b>{item['keywords']}</b>\n\n"
            "Send the new keywords, then /done. /cancel stops reviewing.",
            parse_mode="HTML",
        )
        await callback.answer()
        return

    status = "approved" if action == "approve" else "rejected"
    changed = await db.set_status(media_id, status, callback.from_user.id)
    notice = f"#{media_id} {status}." if changed else f"#{media_id} was already handled."
    await send_next_pending(callback.message.chat.id, notice)
    await callback.answer()


# 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())