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