TelegramBotAdmin/wbb/utils/dbchannel.py

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from uuid import uuid4

from pymongo import ASCENDING, DESCENDING, ReturnDocument

from wbb import BOT_PROFILE_ID, db

postsdb = db.channel_posts
_indexes_ready = False


def now_utc() -> datetime:
    return datetime.now(UTC)


async def ensure_channel_indexes() -> None:
    global _indexes_ready
    if _indexes_ready:
        return
    await postsdb.create_index("post_id", unique=True)
    await postsdb.create_index("telegram_key", unique=True, sparse=True)
    await postsdb.create_index(
        [("bot_id", ASCENDING), ("chat_id", ASCENDING), ("created_at", DESCENDING)]
    )
    await postsdb.create_index(
        [("bot_id", ASCENDING), ("status", ASCENDING), ("publish_at", ASCENDING)]
    )
    _indexes_ready = True


def telegram_key(chat_id: int, message_id: int) -> str:
    return f"{BOT_PROFILE_ID}:{chat_id}:{message_id}"


async def create_post(
    chat_id: int,
    *,
    text: str,
    media_type: str | None,
    file_id: str | None,
    media_filename: str | None,
    media_size: int | None,
    media_mime_type: str | None,
    pin: bool,
    publish_at: datetime | None,
) -> dict:
    await ensure_channel_indexes()
    created = now_utc()
    post = {
        "post_id": uuid4().hex,
        "bot_id": BOT_PROFILE_ID,
        "chat_id": int(chat_id),
        "source": "admin",
        "status": "scheduled" if publish_at else "sending",
        "text": text,
        "media_type": media_type,
        "file_id": file_id,
        "media_filename": media_filename,
        "media_size": media_size,
        "media_mime_type": media_mime_type,
        "pin": bool(pin),
        "publish_at": publish_at,
        "message_id": None,
        "created_at": created,
        "updated_at": created,
    }
    await postsdb.insert_one(post)
    return post


async def get_post(chat_id: int, post_id: str) -> dict | None:
    await ensure_channel_indexes()
    return await postsdb.find_one(
        {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "post_id": post_id}
    )


async def list_posts(chat_id: int, *, page: int, page_size: int) -> tuple[list[dict], int]:
    await ensure_channel_indexes()
    query = {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)}
    total = await postsdb.count_documents(query)
    cursor = (
        postsdb.find(query)
        .sort("created_at", DESCENDING)
        .skip((page - 1) * page_size)
        .limit(page_size)
    )
    return [item async for item in cursor], total


async def update_scheduled_post(chat_id: int, post_id: str, values: dict) -> dict | None:
    await ensure_channel_indexes()
    return await postsdb.find_one_and_update(
        {
            "bot_id": BOT_PROFILE_ID,
            "chat_id": int(chat_id),
            "post_id": post_id,
            "status": "scheduled",
        },
        {"$set": {**values, "updated_at": now_utc()}},
        return_document=ReturnDocument.AFTER,
    )


async def claim_due_posts(limit: int = 20) -> list[dict]:
    await ensure_channel_indexes()
    deadline = now_utc()
    candidates = postsdb.find(
        {
            "bot_id": BOT_PROFILE_ID,
            "status": "scheduled",
            "publish_at": {"$lte": deadline},
        }
    ).sort("publish_at", ASCENDING).limit(limit)
    claimed: list[dict] = []
    async for candidate in candidates:
        item = await postsdb.find_one_and_update(
            {"post_id": candidate["post_id"], "status": "scheduled"},
            {"$set": {"status": "sending", "updated_at": now_utc()}},
            return_document=ReturnDocument.AFTER,
        )
        if item:
            claimed.append(item)
    return claimed


async def mark_stale_sends_uncertain() -> None:
    await ensure_channel_indexes()
    await postsdb.update_many(
        {
            "bot_id": BOT_PROFILE_ID,
            "status": "sending",
            "updated_at": {"$lt": now_utc() - timedelta(minutes=5)},
        },
        {"$set": {"status": "uncertain", "updated_at": now_utc()}},
    )