ctms/backend/app/services/desktop_notification_service.py

from __future__ import annotations

import uuid
from datetime import datetime, timedelta, timezone

from sqlalchemy import and_, exists, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.desktop_notification import (
    DesktopNotificationDelivery,
    DesktopNotificationSubscription,
)
from app.models.notification import Notification
from app.models.study_member import StudyMember
from app.models.user import User

CLAIM_LEASE = timedelta(minutes=5)


async def get_subscription(db: AsyncSession, user_id: uuid.UUID) -> DesktopNotificationSubscription | None:
    return await db.get(DesktopNotificationSubscription, user_id)


async def set_subscription(
    db: AsyncSession,
    user_id: uuid.UUID,
    enabled: bool,
) -> DesktopNotificationSubscription:
    subscription = await get_subscription(db, user_id)
    now = datetime.now(timezone.utc)
    if subscription is None:
        subscription = DesktopNotificationSubscription(
            user_id=user_id,
            enabled=enabled,
            enabled_at=now if enabled else None,
        )
        db.add(subscription)
    elif subscription.enabled != enabled:
        subscription.enabled = enabled
        subscription.enabled_at = now if enabled else None
    await db.commit()
    await db.refresh(subscription)
    return subscription


def _eligible_query(user_id: uuid.UUID, enabled_at: datetime, lease_cutoff: datetime):
    active_membership = exists(
        select(StudyMember.id).where(
            StudyMember.study_id == Notification.study_id,
            StudyMember.user_id == user_id,
            StudyMember.is_active.is_(True),
        )
    )
    return (
        select(Notification, DesktopNotificationDelivery)
        .outerjoin(
            DesktopNotificationDelivery,
            and_(
                DesktopNotificationDelivery.notification_id == Notification.id,
                DesktopNotificationDelivery.user_id == user_id,
            ),
        )
        .where(
            Notification.recipient_id == user_id,
            Notification.created_at >= enabled_at,
            Notification.resolved_at.is_(None),
            Notification.read_at.is_(None),
            active_membership,
            or_(
                DesktopNotificationDelivery.id.is_(None),
                and_(
                    DesktopNotificationDelivery.delivered_at.is_(None),
                    or_(
                        DesktopNotificationDelivery.claimed_at.is_(None),
                        DesktopNotificationDelivery.claimed_at < lease_cutoff,
                    ),
                ),
            ),
        )
        .order_by(Notification.created_at.asc())
        .with_for_update(of=Notification, skip_locked=True)
    )


async def claim_notifications(
    db: AsyncSession,
    user_id: uuid.UUID,
    limit: int,
) -> tuple[uuid.UUID | None, datetime | None, list[Notification]]:
    subscription = await get_subscription(db, user_id)
    if not subscription or not subscription.enabled or not subscription.enabled_at:
        return None, None, []

    # Reconcile every active project before selecting desktop deliveries so permission
    # changes fail closed and time-based reminders do not depend on opening the web Feed.
    user = await db.get(User, user_id)
    if user is None or not user.is_active:
        return None, None, []
    study_ids = (await db.scalars(
        select(StudyMember.study_id).where(
            StudyMember.user_id == user_id,
            StudyMember.is_active.is_(True),
        )
    )).all()
    from app.services.project_reminder_service import sync_project_reminders

    for study_id in set(study_ids):
        await sync_project_reminders(db, study_id, user)

    now = datetime.now(timezone.utc)
    token = uuid.uuid4()
    rows = (
        await db.execute(
            _eligible_query(user_id, subscription.enabled_at, now - CLAIM_LEASE)
            .limit(max(1, min(limit, 50)))
        )
    ).all()
    items: list[Notification] = []
    for notification, delivery in rows:
        if delivery is None:
            delivery = DesktopNotificationDelivery(
                user_id=user_id,
                notification_id=notification.id,
            )
            db.add(delivery)
        delivery.claim_token = token
        delivery.claimed_at = now
        items.append(notification)
    await db.commit()
    if not items:
        return None, None, []
    return token, now + CLAIM_LEASE, items


async def acknowledge_notifications(
    db: AsyncSession,
    user_id: uuid.UUID,
    claim_token: uuid.UUID,
    delivered_ids: list[uuid.UUID],
) -> None:
    if delivered_ids:
        await db.execute(
            update(DesktopNotificationDelivery)
            .where(
                DesktopNotificationDelivery.user_id == user_id,
                DesktopNotificationDelivery.claim_token == claim_token,
                DesktopNotificationDelivery.notification_id.in_(delivered_ids),
            )
            .values(delivered_at=datetime.now(timezone.utc))
        )
    await db.commit()


async def mark_notification_read(
    db: AsyncSession,
    user_id: uuid.UUID,
    notification_id: uuid.UUID,
) -> None:
    item = await db.scalar(select(Notification).where(
        Notification.id == notification_id,
        Notification.recipient_id == user_id,
        Notification.resolved_at.is_(None),
    ))
    if item is None:
        return
    now = datetime.now(timezone.utc)
    item.read_at = item.read_at or now
    if not item.requires_action:
        item.resolved_at = item.resolved_at or now
    await db.commit()