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.distribution import Distribution, DistributionStatus, DistributionTargetType
from app.models.document import Document
from app.models.document_version import DocumentVersion
from app.models.study import Study
from app.models.study_member import StudyMember
from app.schemas.notification import NotificationItem

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):
    role_target_exists = exists(
        select(StudyMember.id).where(
            StudyMember.study_id == Document.trial_id,
            StudyMember.user_id == user_id,
            StudyMember.is_active.is_(True),
            StudyMember.role_in_study == Distribution.target_id,
        )
    )
    target_matches = or_(
        and_(
            Distribution.target_type == DistributionTargetType.USER,
            Distribution.target_id == str(user_id),
        ),
        and_(
            Distribution.target_type == DistributionTargetType.ROLE,
            role_target_exists,
        ),
    )
    return (
        select(
            Distribution,
            Document,
            DocumentVersion,
            Study,
            DesktopNotificationDelivery,
        )
        .join(Document, Distribution.document_id == Document.id)
        .join(DocumentVersion, Distribution.version_id == DocumentVersion.id)
        .join(Study, Document.trial_id == Study.id)
        .outerjoin(
            DesktopNotificationDelivery,
            and_(
                DesktopNotificationDelivery.distribution_id == Distribution.id,
                DesktopNotificationDelivery.user_id == user_id,
            ),
        )
        .where(
            Distribution.status == DistributionStatus.ACTIVE,
            Distribution.created_at >= enabled_at,
            target_matches,
            or_(
                DesktopNotificationDelivery.id.is_(None),
                and_(
                    DesktopNotificationDelivery.delivered_at.is_(None),
                    or_(
                        DesktopNotificationDelivery.claimed_at.is_(None),
                        DesktopNotificationDelivery.claimed_at < lease_cutoff,
                    ),
                ),
            ),
        )
        .order_by(Distribution.created_at.asc())
        .with_for_update(of=Distribution, skip_locked=True)
    )


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

    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[NotificationItem] = []
    for distribution, document, version, study, delivery in rows:
        if delivery is None:
            delivery = DesktopNotificationDelivery(
                user_id=user_id,
                distribution_id=distribution.id,
            )
            db.add(delivery)
        delivery.claim_token = token
        delivery.claimed_at = now
        items.append(
            NotificationItem(
                id=distribution.id,
                document_id=document.id,
                version_id=version.id,
                document_title=document.title,
                document_no=document.doc_no,
                version_no=version.version_no,
                change_summary=version.change_summary,
                effective_at=version.effective_at,
                created_at=distribution.created_at,
                study_id=study.id,
                study_name=study.name,
                delivered_at=delivery.delivered_at,
                read_at=delivery.read_at,
            )
        )
    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.distribution_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,
    distribution_id: uuid.UUID,
) -> None:
    delivery = (
        await db.execute(
            select(DesktopNotificationDelivery).where(
                DesktopNotificationDelivery.user_id == user_id,
                DesktopNotificationDelivery.distribution_id == distribution_id,
            )
        )
    ).scalar_one_or_none()
    if delivery is None:
        delivery = DesktopNotificationDelivery(
            user_id=user_id,
            distribution_id=distribution_id,
        )
        db.add(delivery)
    delivery.read_at = datetime.now(timezone.utc)
    await db.commit()