"""Broadcast campaign, recipient queue and event-log persistence. The recipient table is the durable send queue. A campaign survives a restart because every pending row is still present in the database; no in-memory list is used as the source of truth. """ from __future__ import annotations from datetime import datetime, timedelta from typing import Any from sqlalchemy import func, select, update from app.database.base import utcnow from app.database.models import ( BroadcastAudience, BroadcastCampaign, BroadcastLog, BroadcastMessageType, BroadcastRecipient, BroadcastRecipientStatus, BroadcastStatus, User, ) from app.database.repositories.base import BaseRepository _CAMPAIGN_FIELDS = { "title", "message_type", "text", "parse_mode", "media_file_id", "media_file_name", "media_caption", "buttons", "protect_content", "audience", "active_days", "status", "scheduled_at", "started_at", "finished_at", "total_recipients", "sent_count", "failed_count", "blocked_count", "unknown_count", "skipped_count", "last_error", "correlation_id", } class BroadcastRepository(BaseRepository): async def create_campaign( self, *, created_by: int, title: str | None = None ) -> BroadcastCampaign: async with self.db.session() as session: row = BroadcastCampaign( created_by=created_by, title=title, message_type=BroadcastMessageType.TEXT, status=BroadcastStatus.DRAFT, ) session.add(row) await session.flush() return row async def get_campaign(self, campaign_id: int) -> BroadcastCampaign | None: async with self.db.session() as session: return await session.get(BroadcastCampaign, campaign_id) async def list_campaigns(self, limit: int = 10, offset: int = 0) -> list[BroadcastCampaign]: async with self.db.session() as session: rows = await session.scalars( select(BroadcastCampaign) .order_by(BroadcastCampaign.created_at.desc(), BroadcastCampaign.id.desc()) .limit(limit) .offset(offset) ) return list(rows) async def list_by_status( self, status: BroadcastStatus, limit: int = 10 ) -> list[BroadcastCampaign]: async with self.db.session() as session: rows = await session.scalars( select(BroadcastCampaign) .where(BroadcastCampaign.status == status) .order_by(BroadcastCampaign.id) .limit(limit) ) return list(rows) async def count_campaigns(self) -> int: async with self.db.session() as session: return int(await session.scalar(select(func.count(BroadcastCampaign.id))) or 0) async def update_campaign(self, campaign_id: int, **fields: Any) -> bool: values = {key: value for key, value in fields.items() if key in _CAMPAIGN_FIELDS} if not values: return False values.setdefault("updated_at", utcnow()) async with self.db.session() as session: result = await session.execute( update(BroadcastCampaign) .where(BroadcastCampaign.id == campaign_id) .values(**values) ) return bool(result.rowcount) async def delete_campaign(self, campaign_id: int) -> bool: async with self.db.session() as session: campaign = await session.get(BroadcastCampaign, campaign_id) if campaign is None: return False await session.delete(campaign) return True async def active_campaign(self, *, exclude_id: int | None = None) -> BroadcastCampaign | None: async with self.db.session() as session: statement = select(BroadcastCampaign).where( BroadcastCampaign.status == BroadcastStatus.RUNNING ) if exclude_id is not None: statement = statement.where(BroadcastCampaign.id != exclude_id) return await session.scalar(statement.order_by(BroadcastCampaign.id).limit(1)) async def due_scheduled(self, now: datetime) -> list[BroadcastCampaign]: async with self.db.session() as session: rows = await session.scalars( select(BroadcastCampaign) .where( BroadcastCampaign.status == BroadcastStatus.SCHEDULED, BroadcastCampaign.scheduled_at.is_not(None), BroadcastCampaign.scheduled_at <= now, ) .order_by(BroadcastCampaign.scheduled_at, BroadcastCampaign.id) ) return list(rows) async def audience_count(self, audience: BroadcastAudience, active_days: int | None) -> int: async with self.db.session() as session: statement = select(func.count(User.id)).where( User.is_blocked.is_(False), User.is_broadcast_blocked.is_(False) ) if audience == BroadcastAudience.ACTIVE and active_days: statement = statement.where( User.last_active >= utcnow() - timedelta(days=active_days) ) return int(await session.scalar(statement) or 0) async def add_recipients( self, campaign_id: int, *, audience: BroadcastAudience, active_days: int | None, ) -> int: """Snapshot eligible users into the durable queue. Only users who have started the bot exist in ``users`` at all. Users manually blocked by an administrator or marked as having blocked the bot are excluded. """ async with self.db.session() as session: statement = select(User.id).where( User.is_blocked.is_(False), User.is_broadcast_blocked.is_(False) ) if audience == BroadcastAudience.ACTIVE and active_days: statement = statement.where( User.last_active >= utcnow() - timedelta(days=active_days) ) user_ids = list(await session.scalars(statement.order_by(User.id))) existing = set( await session.scalars( select(BroadcastRecipient.user_id).where( BroadcastRecipient.campaign_id == campaign_id ) ) ) rows = [ BroadcastRecipient( campaign_id=campaign_id, user_id=user_id, status=BroadcastRecipientStatus.PENDING, ) for user_id in user_ids if user_id not in existing ] if rows: session.add_all(rows) await session.flush() return len(rows) async def recipient_counts(self, campaign_id: int) -> dict[str, int]: async with self.db.session() as session: rows = await session.execute( select(BroadcastRecipient.status, func.count(BroadcastRecipient.id)) .where(BroadcastRecipient.campaign_id == campaign_id) .group_by(BroadcastRecipient.status) ) return { (status.value if hasattr(status, "value") else str(status)): int(count) for status, count in rows } async def claim_recipients( self, campaign_id: int, limit: int ) -> list[tuple[int, int]]: """Atomically move up to ``limit`` pending rows to SENDING.""" async with self.db.session() as session: candidates = list( await session.execute( select(BroadcastRecipient.id, User.telegram_id) .join(User, User.id == BroadcastRecipient.user_id) .where( BroadcastRecipient.campaign_id == campaign_id, BroadcastRecipient.status == BroadcastRecipientStatus.PENDING, ) .order_by(BroadcastRecipient.id) .limit(limit) ) ) claimed: list[tuple[int, int]] = [] for recipient_id, telegram_id in candidates: result = await session.execute( update(BroadcastRecipient) .where( BroadcastRecipient.id == recipient_id, BroadcastRecipient.status == BroadcastRecipientStatus.PENDING, ) .values( status=BroadcastRecipientStatus.SENDING, attempts=BroadcastRecipient.attempts + 1, ) ) if result.rowcount: claimed.append((int(recipient_id), int(telegram_id))) return claimed async def mark_recipient( self, recipient_id: int, status: BroadcastRecipientStatus, *, telegram_message_id: int | None = None, error: str | None = None, ) -> bool: values: dict[str, Any] = { "status": status, "telegram_message_id": telegram_message_id, "error_message": (error or None) and error[:1000], "updated_at": utcnow(), } if status == BroadcastRecipientStatus.SENT: values["sent_at"] = utcnow() async with self.db.session() as session: result = await session.execute( update(BroadcastRecipient) .where(BroadcastRecipient.id == recipient_id) .values(**values) ) return bool(result.rowcount) async def recover_sending(self) -> int: """Crash recovery: never blindly resend an in-flight broadcast.""" async with self.db.session() as session: result = await session.execute( update(BroadcastRecipient) .where(BroadcastRecipient.status == BroadcastRecipientStatus.SENDING) .values( status=BroadcastRecipientStatus.UNKNOWN, error_message="نتیجه ارسال پس از راه‌اندازی مجدد نامشخص است.", updated_at=utcnow(), ) ) return int(result.rowcount or 0) async def reset_failed(self, campaign_id: int) -> int: async with self.db.session() as session: result = await session.execute( update(BroadcastRecipient) .where( BroadcastRecipient.campaign_id == campaign_id, BroadcastRecipient.status == BroadcastRecipientStatus.FAILED, ) .values( status=BroadcastRecipientStatus.PENDING, attempts=0, error_message=None, updated_at=utcnow(), ) ) return int(result.rowcount or 0) async def skip_pending(self, campaign_id: int) -> int: async with self.db.session() as session: result = await session.execute( update(BroadcastRecipient) .where( BroadcastRecipient.campaign_id == campaign_id, BroadcastRecipient.status.in_( [ BroadcastRecipientStatus.PENDING, BroadcastRecipientStatus.SENDING, ] ), ) .values( status=BroadcastRecipientStatus.SKIPPED, error_message="کمپین لغو شد.", updated_at=utcnow(), ) ) return int(result.rowcount or 0) async def log_event( self, campaign_id: int, event: str, *, details: dict[str, Any] | None = None, actor_id: int | None = None, ) -> None: async with self.db.session() as session: session.add( BroadcastLog( campaign_id=campaign_id, event=event[:64], details=details or {}, actor_id=actor_id, ) ) async def list_logs(self, campaign_id: int, limit: int = 20) -> list[BroadcastLog]: async with self.db.session() as session: rows = await session.scalars( select(BroadcastLog) .where(BroadcastLog.campaign_id == campaign_id) .order_by(BroadcastLog.created_at.desc(), BroadcastLog.id.desc()) .limit(limit) ) return list(rows) async def totals(self) -> tuple[int, int, int, int]: """Return (campaigns, sent, failed, blocked) for the dashboard.""" async with self.db.session() as session: row = ( await session.execute( select( func.count(BroadcastCampaign.id), func.coalesce(func.sum(BroadcastCampaign.sent_count), 0), func.coalesce(func.sum(BroadcastCampaign.failed_count), 0), func.coalesce(func.sum(BroadcastCampaign.blocked_count), 0), ) ) ).one() return int(row[0]), int(row[1]), int(row[2]), int(row[3])