import logging from datetime import datetime, timezone, timedelta from sqlalchemy import select, func, update, case, or_ from sqlalchemy.ext.asyncio import AsyncSession from app.models.notification import Notification from app.models.notification_read import NotificationRead from app.utils.id_gen import generate_id logger = logging.getLogger("videogen") CST = timezone(timedelta(hours=8)) def _to_local_str(dt): if dt is None: return None if dt.tzinfo and dt.tzinfo.utcoffset(None) == timedelta(0): dt = dt.astimezone(CST) return dt.replace(tzinfo=None).isoformat() async def create_notification( db: AsyncSession, user_id: str | None, title: str, content: str, notif_type: str = "system", related_id: str | None = None, push_ws: bool = True, ) -> Notification: """Create a notification. user_id=None means broadcast.""" notif = Notification( id=generate_id(), user_id=user_id, title=title, content=content, type=notif_type, related_id=related_id, ) db.add(notif) await db.flush() # Push via WebSocket if available if push_ws: try: from app.utils.redis import get_redis redis = get_redis() if redis: import json payload = json.dumps( { "event": "notification", "data": { "id": notif.id, "title": notif.title, "content": notif.content, "type": notif.type, "user_id": notif.user_id, "created_at": _to_local_str(notif.created_at), }, } ) if notif.user_id: await redis.publish(f"user:{notif.user_id}:notifications", payload) else: await redis.publish("broadcast:notifications", payload) except Exception: pass return notif async def get_notifications( db: AsyncSession, user_id: str, page: int = 1, page_size: int = 20 ) -> tuple[list[dict], int]: """Get notifications for a user (personal + broadcast), with per-user read state.""" query = ( select( Notification, case( (Notification.user_id.is_(None), NotificationRead.id.is_not(None)), else_=Notification.is_read, ).label("user_is_read"), ) .outerjoin( NotificationRead, (Notification.id == NotificationRead.notification_id) & (NotificationRead.user_id == user_id), ) .where( or_( Notification.user_id == user_id, Notification.user_id.is_(None), ) ) .order_by(Notification.created_at.desc()) ) count_query = ( select(func.count(Notification.id)) .where( or_( Notification.user_id == user_id, Notification.user_id.is_(None), ) ) ) total = (await db.execute(count_query)).scalar() or 0 result = await db.execute(query.offset((page - 1) * page_size).limit(page_size)) rows = result.all() items = [] for notif, user_is_read in rows: items.append({ "id": notif.id, "user_id": notif.user_id, "title": notif.title, "content": notif.content, "type": notif.type, "is_read": bool(user_is_read), "related_id": notif.related_id, "created_at": notif.created_at, }) return items, total async def mark_read(db: AsyncSession, notification_id: str, user_id: str) -> None: """Mark a notification as read for a specific user.""" result = await db.execute( select(Notification).where(Notification.id == notification_id).limit(1) ) notif = result.scalar_one_or_none() if not notif: return if notif.user_id is not None: # Personal notification: update the row directly if notif.user_id == user_id: notif.is_read = True await db.flush() else: # Broadcast notification: use notification_reads table existing = await db.execute( select(NotificationRead).where( NotificationRead.notification_id == notification_id, NotificationRead.user_id == user_id, ) .limit(1) ) if not existing.scalar_one_or_none(): db.add(NotificationRead( id=generate_id(), notification_id=notification_id, user_id=user_id, )) await db.flush() async def mark_all_read(db: AsyncSession, user_id: str) -> None: """Mark all notifications as read for a user.""" # Mark personal notifications await db.execute( update(Notification) .where( Notification.user_id == user_id, Notification.is_read == False, ) .values(is_read=True) ) # Mark all broadcast notifications broadcasts = await db.execute( select(Notification.id) .where(Notification.user_id.is_(None)) ) broadcast_ids = [row[0] for row in broadcasts.all()] if broadcast_ids: existing_reads = await db.execute( select(NotificationRead.notification_id).where( NotificationRead.notification_id.in_(broadcast_ids), NotificationRead.user_id == user_id, ) ) already_read = {row[0] for row in existing_reads.all()} for bid in broadcast_ids: if bid not in already_read: db.add(NotificationRead( id=generate_id(), notification_id=bid, user_id=user_id, )) await db.flush() async def get_unread_count(db: AsyncSession, user_id: str) -> int: """Get unread notification count for a user.""" # Count personal unread personal_count = ( await db.execute( select(func.count(Notification.id)).where( Notification.user_id == user_id, Notification.is_read == False, ) ) ).scalar() or 0 # Count broadcast unread (not in notification_reads) read_broadcast_ids = ( await db.execute( select(NotificationRead.notification_id).where( NotificationRead.user_id == user_id, ) ) ).all() read_ids = {row[0] for row in read_broadcast_ids} broadcast_query = select(func.count(Notification.id)).where( Notification.user_id.is_(None), ) if read_ids: broadcast_query = broadcast_query.where(~Notification.id.in_(read_ids)) broadcast_count = (await db.execute(broadcast_query)).scalar() or 0 return personal_count + broadcast_count async def get_notification_read_users( db: AsyncSession, notification_id: str, page: int = 1, page_size: int = 50 ) -> tuple[list[dict], int]: """Get list of users who have read a specific notification.""" from app.models.user import User query = ( select(NotificationRead, User.username) .join(User, NotificationRead.user_id == User.id) .where(NotificationRead.notification_id == notification_id) .order_by(NotificationRead.created_at.desc()) ) count_query = ( select(func.count(NotificationRead.id)) .where(NotificationRead.notification_id == notification_id) ) total = (await db.execute(count_query)).scalar() or 0 result = await db.execute(query.offset((page - 1) * page_size).limit(page_size)) rows = result.all() items = [] for read_record, username in rows: items.append({ "user_id": read_record.user_id, "username": username, "read_at": read_record.created_at, }) return items, total