FabledScribe/src/fabledassistant/services/push.py

"""Browser push notification service using VAPID/pywebpush."""
import asyncio
import base64
import json
import logging
from datetime import datetime, timezone
from functools import lru_cache
from pathlib import Path

from sqlalchemy import delete, select

from fabledassistant.config import Config
from fabledassistant.models import async_session
from fabledassistant.models.push_subscription import PushSubscription

logger = logging.getLogger(__name__)

# Persisted alongside other app data so keys survive container restarts.
_VAPID_KEYS_FILE = Path(Config.IMAGE_CACHE_DIR).parent / "vapid_keys.json"


@lru_cache(maxsize=1)
def _get_webpush():
    """Lazy import to avoid startup errors if pywebpush is not installed."""
    try:
        from pywebpush import webpush
        return webpush
    except ImportError:
        return None


def vapid_enabled() -> bool:
    return bool(Config.VAPID_PRIVATE_KEY and Config.VAPID_PUBLIC_KEY)


def ensure_vapid_keys() -> None:
    """Load or auto-generate VAPID keys, storing them in the data volume.

    Called once at startup.  If VAPID_PRIVATE_KEY / VAPID_PUBLIC_KEY env vars
    are already set they take precedence and nothing is written to disk.
    """
    if vapid_enabled():
        logger.info("VAPID keys loaded from environment variables")
        return

    # Try to load previously generated keys from the data volume.
    if _VAPID_KEYS_FILE.exists():
        try:
            data = json.loads(_VAPID_KEYS_FILE.read_text())
            Config.VAPID_PRIVATE_KEY = data["private_key"]
            Config.VAPID_PUBLIC_KEY = data["public_key"]
            logger.info("VAPID keys loaded from %s", _VAPID_KEYS_FILE)
            return
        except Exception:
            logger.warning("Failed to load VAPID keys from %s — regenerating", _VAPID_KEYS_FILE, exc_info=True)

    # Generate a fresh key pair.
    try:
        from cryptography.hazmat.primitives import serialization
        from py_vapid import Vapid01

        v = Vapid01()
        v.generate_keys()

        private_pem = v.private_key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.TraditionalOpenSSL,
            serialization.NoEncryption(),
        ).decode()

        pub_bytes = v.public_key.public_bytes(
            serialization.Encoding.X962,
            serialization.PublicFormat.UncompressedPoint,
        )
        public_b64 = base64.urlsafe_b64encode(pub_bytes).rstrip(b"=").decode()

        # Persist so they survive container restarts.
        _VAPID_KEYS_FILE.parent.mkdir(parents=True, exist_ok=True)
        _VAPID_KEYS_FILE.write_text(json.dumps({"private_key": private_pem, "public_key": public_b64}))

        Config.VAPID_PRIVATE_KEY = private_pem
        Config.VAPID_PUBLIC_KEY = public_b64
        logger.info("Generated new VAPID keys and saved to %s", _VAPID_KEYS_FILE)
    except Exception:
        logger.warning("Failed to generate VAPID keys — push notifications will be unavailable", exc_info=True)


async def save_subscription(user_id: int, subscription_json: dict, user_agent: str | None = None) -> PushSubscription:
    """Upsert a push subscription by endpoint."""
    endpoint = subscription_json.get("endpoint", "")
    keys = subscription_json.get("keys", {})
    p256dh = keys.get("p256dh", "")
    auth = keys.get("auth", "")

    async with async_session() as session:
        result = await session.execute(
            select(PushSubscription).where(
                PushSubscription.user_id == user_id,
                PushSubscription.endpoint == endpoint,
            )
        )
        existing = result.scalars().first()
        if existing:
            existing.p256dh = p256dh
            existing.auth = auth
            existing.user_agent = user_agent
            existing.last_used = datetime.now(timezone.utc)
            await session.commit()
            await session.refresh(existing)
            return existing
        else:
            sub = PushSubscription(
                user_id=user_id,
                endpoint=endpoint,
                p256dh=p256dh,
                auth=auth,
                user_agent=user_agent,
            )
            session.add(sub)
            await session.commit()
            await session.refresh(sub)
            return sub


async def delete_subscription(user_id: int, endpoint: str) -> None:
    async with async_session() as session:
        await session.execute(
            delete(PushSubscription).where(
                PushSubscription.user_id == user_id,
                PushSubscription.endpoint == endpoint,
            )
        )
        await session.commit()


async def _remove_expired_subscription(user_id: int, endpoint: str) -> None:
    """Remove a subscription that returned 410 Gone."""
    try:
        await delete_subscription(user_id, endpoint)
        logger.info("Removed expired push subscription for user %d", user_id)
    except Exception:
        logger.warning("Failed to remove expired subscription", exc_info=True)


async def send_push_notification(
    user_id: int,
    title: str,
    body: str,
    url: str = "/",
) -> None:
    """Send a push notification to all subscriptions for a user.

    Fire-and-forget — wrap in asyncio.create_task() at call site.
    """
    if not vapid_enabled():
        logger.debug("VAPID not configured, skipping push notification")
        return

    webpush = _get_webpush()
    if webpush is None:
        logger.warning("pywebpush not installed, cannot send push notifications")
        return

    async with async_session() as session:
        result = await session.execute(
            select(PushSubscription).where(PushSubscription.user_id == user_id)
        )
        subscriptions = list(result.scalars().all())

    if not subscriptions:
        return

    payload = json.dumps({"title": title, "body": body, "url": url})
    vapid_claims = {
        "sub": Config.VAPID_CLAIMS_SUB,
    }

    def _send_sync(sub: PushSubscription) -> tuple[int, str]:
        """Synchronous send — runs in executor."""
        try:
            subscription_info = {
                "endpoint": sub.endpoint,
                "keys": {"p256dh": sub.p256dh, "auth": sub.auth},
            }
            response = webpush(
                subscription_info=subscription_info,
                data=payload,
                vapid_private_key=Config.VAPID_PRIVATE_KEY,
                vapid_claims=vapid_claims,
            )
            return response.status_code, sub.endpoint
        except Exception as e:
            logger.error("Push send failed for sub %d: %s", sub.id, e, exc_info=True)
            return 0, sub.endpoint

    loop = asyncio.get_event_loop()
    tasks = [loop.run_in_executor(None, _send_sync, sub) for sub in subscriptions]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    for sub, result in zip(subscriptions, results):
        if isinstance(result, Exception):
            logger.error("Push gather exception for user %d: %s", user_id, result, exc_info=result)
            continue
        status_code, endpoint = result
        if status_code in (200, 201):
            logger.info("Push notification sent to sub %d (status %d)", sub.id, status_code)
        elif status_code == 410:
            asyncio.create_task(_remove_expired_subscription(user_id, endpoint))
        elif status_code:
            logger.error("Push returned unexpected status %d for user %d sub %d", status_code, user_id, sub.id)