import asyncio
import logging
import os
import time
import uuid

from aiogram import Bot, Dispatcher, F
from aiogram.exceptions import TelegramAPIError
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,
    InlineQueryResultsButton,
    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()
}

def _env_flag(name: str) -> bool:
    return os.environ.get(name, "").strip().lower() in {"1", "true", "yes"}


# 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 = _env_flag("ACCEPT_SUBMISSIONS")

# Public mode swaps the ALLOWED_USERS allowlist for "is this user in the chats
# I require?". With no REQUIRED_CHATS the bot is simply open to everyone.
PUBLIC_MODE = _env_flag("PUBLIC_MODE")


def _parse_chat(raw: str) -> int | str:
    """Telegram takes either @username or a numeric id."""
    raw = raw.strip()
    try:
        return int(raw)
    except ValueError:
        return raw


REQUIRED_CHATS: list[int | str] = [
    _parse_chat(c) for c in os.environ.get("REQUIRED_CHATS", "").split(",") if c.strip()
]
# "all" (default): must be in every listed chat. "any": one is enough.
REQUIRE_MODE = os.environ.get("REQUIRE_MODE", "all").strip().lower()

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="vault", description="Retrieve a saved post"),
    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 ADMIN_USERS:
        return user_id in ADMIN_USERS
    # An unconfigured private bot treats its users as admins, which is the
    # original single-owner behaviour. That fallback must not survive into
    # public mode, where it would hand admin rights to anyone who found the
    # bot. check_admin_config() flags the missing setting at startup.
    if PUBLIC_MODE:
        return False
    return not ALLOWED_USERS or user_id in ALLOWED_USERS


# Inline queries fire on every keystroke, so an uncached membership check would
# mean several API calls per typed character. Denials expire quickly so that a
# user who has just joined is not left waiting for the full window.
MEMBERSHIP_TTL_ALLOWED = 300.0
MEMBERSHIP_TTL_DENIED = 30.0

ACCEPTED_MEMBER_STATUSES = {"creator", "administrator", "member"}

_membership_cache: dict[tuple[int, int | str], tuple[bool, float]] = {}


async def _fetch_membership(user_id: int, chat: int | str) -> bool:
    try:
        member = await bot.get_chat_member(chat, user_id)
    except TelegramAPIError:
        # Usually the bot is not an admin in that chat, or the user is unknown
        # to it. Fail closed; check_required_chats() reports setup problems at
        # startup so this does not fail closed silently forever.
        log.warning("membership check failed for user %s in chat %s", user_id, chat, exc_info=True)
        return False

    if member.status in ACCEPTED_MEMBER_STATUSES:
        return True
    # A restricted user is still a member unless they have actually left.
    if member.status == "restricted":
        return bool(getattr(member, "is_member", False))
    return False


async def is_member_of(user_id: int, chat: int | str) -> bool:
    key = (user_id, chat)
    cached = _membership_cache.get(key)
    now = time.monotonic()
    if cached is not None and cached[1] > now:
        return cached[0]

    result = await _fetch_membership(user_id, chat)
    ttl = MEMBERSHIP_TTL_ALLOWED if result else MEMBERSHIP_TTL_DENIED
    _membership_cache[key] = (result, now + ttl)
    return result


async def passes_membership(user_id: int) -> bool:
    if not REQUIRED_CHATS:
        return True
    if REQUIRE_MODE == "any":
        for chat in REQUIRED_CHATS:
            if await is_member_of(user_id, chat):
                return True
        return False
    for chat in REQUIRED_CHATS:
        if not await is_member_of(user_id, chat):
            return False
    return True


async def has_access(user_id: int) -> bool:
    """Public mode replaces the allowlist with the membership requirement."""
    if is_admin(user_id):
        return True
    if PUBLIC_MODE:
        return await passes_membership(user_id)
    return is_allowed(user_id)


# Resolved once at startup so the join prompt can show real links.
REQUIRED_CHAT_LINKS: list[tuple[str, str | None]] = []


def check_admin_config() -> None:
    if PUBLIC_MODE and not ADMIN_USERS:
        log.error(
            "PUBLIC_MODE is on but ADMIN_USERS is empty — nobody can moderate, "
            "add or delete. Set ADMIN_USERS to your Telegram user ID."
        )


async def check_required_chats() -> None:
    """Verify the bot can actually see each required chat, and remember how to
    link to it. A misconfigured chat would otherwise lock out every user with
    nothing in the logs to explain why."""
    REQUIRED_CHAT_LINKS.clear()
    for chat in REQUIRED_CHATS:
        try:
            info = await bot.get_chat(chat)
        except TelegramAPIError:
            log.error(
                "REQUIRED_CHATS: cannot read %s — add the bot to that chat as an "
                "admin, or users will not be able to pass the membership check.",
                chat, exc_info=True,
            )
            REQUIRED_CHAT_LINKS.append((str(chat), None))
            continue
        url = f"https://t.me/{info.username}" if info.username else info.invite_link
        REQUIRED_CHAT_LINKS.append((info.title or str(chat), url))
        if url is None:
            log.warning("REQUIRED_CHATS: no public link for %s; users will only see its name", chat)


def join_prompt() -> str:
    if not REQUIRED_CHAT_LINKS:
        return "You do not have access to this bot."
    joiner = "any one of" if REQUIRE_MODE == "any" else "all of"
    lines = [f"To use this bot you need to join {joiner} these:"]
    for title, url in REQUIRED_CHAT_LINKS:
        lines.append(f"• <a href=\"{url}\">{title}</a>" if url else f"• {title}")
    return "\n".join(lines)


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


class VaultSave(StatesGroup):
    waiting_for_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"
            "/vault [search] — retrieve a whole post you saved\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\n"
        "/vault [search] — retrieve a whole post you saved"
    )


async def deny(msg: Message) -> None:
    """A public bot tells the user how to get in; a private one stays silent,
    as it always has, rather than confirming to strangers that it exists."""
    if PUBLIC_MODE:
        await msg.answer(join_prompt(), parse_mode="HTML")


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

@dp.message(Command("start"))
async def cmd_start(msg: Message):
    if not await has_access(msg.from_user.id):
        await deny(msg)
        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 await has_access(msg.from_user.id):
        await deny(msg)
        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 await has_access(msg.from_user.id):
        return
    await state.clear()
    await msg.answer("Cancelled.")


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

async def start_keyword_prompt(chat_id: int, state: FSMContext, items: list[dict],
                               status: str = "approved") -> None:
    """Begin collecting keywords for one or more media items. Everything in
    `items` ends up sharing the keywords, which is what makes extracting a
    whole album a single interaction."""
    await state.set_state(AddMedia.waiting_for_keywords)
    await state.update_data(items=items, keywords=[], status=status)
    count = f"{len(items)} items" if len(items) > 1 else "it"
    await bot.send_message(
        chat_id,
        f"Got {count}! Now send keywords — one per message or comma-separated.\n"
        "Send /done when finished.",
        parse_mode="HTML",
    )


async def _receive_media(msg: Message, state: FSMContext, file_id: str, file_unique_id: str,
                         media_type: str, status: str = "approved"):
    item = {"file_id": file_id, "file_unique_id": file_unique_id, "media_type": media_type}
    await start_keyword_prompt(msg.chat.id, state, [item], status)


# 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), F.forward_origin.is_(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 await has_access(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")
    items = data.get("items", [])
    for item in items:
        await db.add_media(
            msg.from_user.id,
            item["file_id"],
            item["file_unique_id"],
            item["media_type"],
            keywords,
            status=status,
        )
    await state.clear()
    saved = f"{len(items)} items" if len(items) > 1 else "it"
    if status == "pending":
        await msg.answer(
            f"Submitted {saved} 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 {saved} with keywords: <b>{', '.join(keywords)}</b>",
                         parse_mode="HTML")


async def accumulate_keywords(msg: Message, state: FSMContext) -> None:
    """Shared by adding, editing and vault saves: merge in whatever was typed
    and echo the running list back."""
    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, F.text)
async def on_keyword(msg: Message, state: FSMContext):
    await accumulate_keywords(msg, state)


@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 await has_access(msg.from_user.id):
        await deny(msg)
        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 await has_access(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):
    await accumulate_keywords(msg, state)


# ── forwards: save whole, or extract the media ───────────────────────────────
#
# A forwarded album arrives as several separate messages sharing a
# media_group_id, with no signal that the last one has landed, so they are
# buffered briefly and treated as one post.

ALBUM_DEBOUNCE = 1.0
PENDING_FORWARD_LIMIT = 200

_album_buffers: dict[str, dict] = {}
_pending_forwards: dict[str, dict] = {}


def _new_forward_record(msg: Message) -> dict:
    return {
        "user_id": msg.from_user.id,
        "chat_id": msg.chat.id,
        "message_ids": [],
        "caption": "",
        "media": [],
    }


def _absorb(record: dict, msg: Message) -> None:
    record["message_ids"].append(msg.message_id)
    # In an album only one message carries the caption.
    text = msg.caption or msg.text or ""
    if text and not record["caption"]:
        record["caption"] = text
    found = extract_media(msg)
    if found is not None:
        media_type, file_id, file_unique_id = found
        record["media"].append(
            {"media_type": media_type, "file_id": file_id, "file_unique_id": file_unique_id}
        )


async def prompt_forward_choice(record: dict) -> None:
    token = uuid.uuid4().hex[:12]
    # Bounded so a long-running bot cannot accumulate abandoned prompts.
    while len(_pending_forwards) >= PENDING_FORWARD_LIMIT:
        _pending_forwards.pop(next(iter(_pending_forwards)))
    _pending_forwards[token] = record

    record["message_ids"].sort()
    parts = [f"{len(record['message_ids'])} message(s)"]
    if record["media"]:
        parts.append(f"{len(record['media'])} media")
    summary = ", ".join(parts)

    row = [InlineKeyboardButton(text="📥 Save whole post", callback_data=f"fwd:save:{token}")]
    if record["media"]:
        row.append(InlineKeyboardButton(text="✂️ Extract media", callback_data=f"fwd:extract:{token}"))

    await bot.send_message(
        record["chat_id"],
        f"Forwarded post received ({summary}).\n\n"
        "<b>Save whole post</b> keeps it intact in your vault — retrievable only "
        "inside this bot, with /vault.\n"
        "<b>Extract media</b> adds the media to your searchable library as usual.",
        parse_mode="HTML",
        reply_markup=InlineKeyboardMarkup(inline_keyboard=[row]),
    )


async def _flush_album(group_id: str) -> None:
    """Wait until an album has stopped arriving, then prompt once for all of it."""
    while True:
        record = _album_buffers.get(group_id)
        if record is None:
            return
        remaining = record["deadline"] - time.monotonic()
        if remaining <= 0:
            break
        await asyncio.sleep(remaining)

    record = _album_buffers.pop(group_id, None)
    if record is not None:
        await prompt_forward_choice(record)


@dp.message(F.forward_origin, StateFilter(None))
async def on_forwarded(msg: Message, state: FSMContext):
    if not await has_access(msg.from_user.id):
        await deny(msg)
        return

    group_id = msg.media_group_id
    if group_id is None:
        record = _new_forward_record(msg)
        _absorb(record, msg)
        await prompt_forward_choice(record)
        return

    record = _album_buffers.get(group_id)
    if record is None:
        record = _new_forward_record(msg)
        _album_buffers[group_id] = record
        record["deadline"] = time.monotonic() + ALBUM_DEBOUNCE
        asyncio.create_task(_flush_album(group_id))
    _absorb(record, msg)
    record["deadline"] = time.monotonic() + ALBUM_DEBOUNCE


@dp.callback_query(F.data.startswith("fwd:"))
async def on_forward_choice(callback: CallbackQuery, state: FSMContext):
    _, action, token = callback.data.split(":", 2)
    record = _pending_forwards.pop(token, None)

    if record is None:
        await callback.answer("That prompt expired — forward the post again.", show_alert=True)
        return
    if record["user_id"] != callback.from_user.id:
        await callback.answer("Not yours.", show_alert=True)
        return

    await callback.message.edit_reply_markup(reply_markup=None)

    if action == "extract":
        if not await has_access(callback.from_user.id):
            await callback.answer("Not allowed.", show_alert=True)
            return
        status = "approved" if is_admin(callback.from_user.id) else "pending"
        await start_keyword_prompt(record["chat_id"], state, record["media"], status)
        await callback.answer()
        return

    await state.set_state(VaultSave.waiting_for_keywords)
    await state.update_data(
        chat_id=record["chat_id"],
        message_ids=record["message_ids"],
        caption=record["caption"],
        keywords=[],
    )
    hint = "Send keywords, then /done." if record["caption"] else \
        "Send keywords, then /done. This post has no caption, so keywords are how you will find it."
    await callback.message.answer(hint)
    await callback.answer()


@dp.message(VaultSave.waiting_for_keywords, Command("done"))
async def cmd_vault_done(msg: Message, state: FSMContext):
    data = await state.get_data()
    keywords: list[str] = data.get("keywords", [])
    caption: str = data.get("caption", "")

    if not keywords and not caption:
        await msg.answer("Add at least one keyword — there is no caption to find this by.")
        return

    item_id = await db.add_vault_item(
        msg.from_user.id, data["chat_id"], data["message_ids"], caption, keywords
    )
    await state.clear()
    await msg.answer(
        f"Saved to your vault as <b>#{item_id}</b>.\n\n"
        "Vault posts come back whole, exactly as forwarded — but only here in "
        "this chat, with /vault. Inline search does not reach them.",
        parse_mode="HTML",
    )


@dp.message(VaultSave.waiting_for_keywords, F.text)
async def on_vault_keyword(msg: Message, state: FSMContext):
    await accumulate_keywords(msg, state)


# ── /vault ───────────────────────────────────────────────────────────────────

def vault_label(item: dict) -> str:
    text = item["keywords"].replace(",", ", ") or item["caption"]
    text = " ".join(text.split())
    if len(text) > 40:
        text = text[:39] + "…"
    count = len(db.unpack_message_ids(item["message_ids"]))
    suffix = f" ({count})" if count > 1 else ""
    return f"#{item['id']} {text or 'untitled'}{suffix}"


@dp.message(Command("vault"))
async def cmd_vault(msg: Message):
    if not await has_access(msg.from_user.id):
        await deny(msg)
        return

    parts = msg.text.split(maxsplit=1)
    query = parts[1].strip() if len(parts) > 1 else ""
    items = await db.search_vault(msg.from_user.id, query)

    if not items:
        if query:
            await msg.answer(f"Nothing in your vault matching '{query}'.")
        else:
            await msg.answer(
                "Your vault is empty. Forward a post to me and choose "
                "<b>Save whole post</b> to put it here.",
                parse_mode="HTML",
            )
        return

    rows = [
        [
            InlineKeyboardButton(text=vault_label(item), callback_data=f"vault:get:{item['id']}"),
            InlineKeyboardButton(text="🗑", callback_data=f"vault:del:{item['id']}"),
        ]
        for item in items
    ]
    total = await db.count_vault(msg.from_user.id)
    header = f"Showing {len(items)} of {total} saved posts." if not query else \
        f"{len(items)} match '{query}'."
    await msg.answer(header, reply_markup=InlineKeyboardMarkup(inline_keyboard=rows))


@dp.callback_query(F.data.startswith("vault:"))
async def on_vault_action(callback: CallbackQuery):
    _, action, raw_id = callback.data.split(":", 2)
    # Scoped by owner, so a guessed id cannot reach someone else's vault.
    item = await db.get_vault_item(callback.from_user.id, int(raw_id))
    if item is None:
        await callback.answer("Not found.", show_alert=True)
        return

    if action == "del":
        await db.delete_vault_item(callback.from_user.id, item["id"])
        await callback.answer("Removed from vault.")
        await callback.message.answer(f"Removed #{item['id']} from your vault.")
        return

    message_ids = db.unpack_message_ids(item["message_ids"])
    try:
        # forwardMessages keeps an album together and preserves the original
        # "forwarded from" header. Telegram requires ascending ids.
        await bot.forward_messages(
            chat_id=callback.message.chat.id,
            from_chat_id=item["chat_id"],
            message_ids=sorted(message_ids),
        )
        await callback.answer()
    except TelegramAPIError:
        log.exception("could not forward vault item %s", item["id"])
        await callback.answer("Could not retrieve that post.", show_alert=True)
        await callback.message.answer(
            f"#{item['id']} could not be retrieved — the original messages may have "
            "been deleted from this chat. Forward the post again to re-save it."
        )


# ── /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 await has_access(query.from_user.id):
        # A bar above the results beats silently returning nothing: tapping it
        # opens the bot in private chat, where /start explains what to join.
        button = None
        if PUBLIC_MODE and REQUIRED_CHATS:
            button = InlineQueryResultsButton(
                text="Tap to see what you need to join", start_parameter="join"
            )
        await query.answer([], cache_time=1, is_personal=True, button=button)
        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()
    check_admin_config()
    await check_required_chats()
    await register_commands()
    log.info("Bot starting…")
    await dp.start_polling(bot)


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