from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from pydantic import BaseModel from typing import Optional from datetime import datetime, timezone from app.db.base import get_db from app.models.notification import Notification from app.api.auth import get_current_user from app.core.security import get_current_user_sse from app.models.user import User, UserRole from app.core.notifications import notification_manager from fastapi.responses import StreamingResponse import asyncio import json router = APIRouter(prefix="/notifications", tags=["Notifications"]) def _check_admin(user_id: int, db: Session): user = db.query(User).filter(User.id == user_id).first() if not user or user.role != UserRole.ADMIN: raise HTTPException(status_code=403, detail="Admin access required") return user # ── Schemas ─────────────────────────────────────────────────── class NotificationCreate(BaseModel): type: str # checkout_shipping | checkout_payment | checkout_complete title: str message: Optional[str] = None data: Optional[dict] = None class NotificationOut(BaseModel): id: int type: str title: str message: Optional[str] data: Optional[dict] is_read: bool created_at: datetime class Config: from_attributes = True # ── POST /notifications — create (no auth, called from checkout) ── @router.post("/") async def create_notification(payload: NotificationCreate, db: Session = Depends(get_db)): notif = Notification( type=payload.type, title=payload.title, message=payload.message, data=payload.data, is_read=False, created_at=datetime.now(timezone.utc), ) db.add(notif) db.commit() db.refresh(notif) # Broadcast to active SSE clients await notification_manager.broadcast( NotificationOut.model_validate(notif).model_dump() ) return { "isSuccess": True, "value": {"id": notif.id}, "statusCode": 201, } # ── GET /notifications — list for admin ── @router.get("/") async def list_notifications( skip: int = 0, limit: int = 50, unread_only: bool = False, current_user_id: int = Depends(get_current_user), db: Session = Depends(get_db), ): _check_admin(current_user_id, db) query = db.query(Notification) if unread_only: query = query.filter(Notification.is_read == False) # noqa: E712 query = query.order_by(Notification.created_at.desc()) total = query.count() notifications = query.offset(skip).limit(limit).all() return { "isSuccess": True, "value": { "notifications": [ NotificationOut.model_validate(n).model_dump() for n in notifications ], "total": total, "unread_count": db.query(Notification).filter(Notification.is_read == False).count(), # noqa: E712 }, "statusCode": 200, } # ── PUT /notifications/{id}/read — mark single as read ── @router.put("/{notification_id}/read") async def mark_read( notification_id: int, current_user_id: int = Depends(get_current_user), db: Session = Depends(get_db), ): _check_admin(current_user_id, db) notif = db.query(Notification).filter(Notification.id == notification_id).first() if not notif: raise HTTPException(status_code=404, detail="Notification not found") notif.is_read = True db.commit() return {"isSuccess": True, "value": None, "statusCode": 200} # ── PUT /notifications/read-all — mark all as read ── @router.put("/read-all") async def mark_all_read( current_user_id: int = Depends(get_current_user), db: Session = Depends(get_db), ): _check_admin(current_user_id, db) db.query(Notification).filter(Notification.is_read == False).update( # noqa: E712 {"is_read": True} ) db.commit() return {"isSuccess": True, "value": None, "statusCode": 200} # ── GET /notifications/stream — SSE stream for admin ── @router.get("/stream") async def stream_notifications( current_user_id: int = Depends(get_current_user_sse), db: Session = Depends(get_db), ): """ SSE endpoint for real-time notifications. Includes a heartbeat to keep the connection alive on Hugging Face / Vercel. """ _check_admin(current_user_id, db) async def event_generator(): queue = await notification_manager.subscribe(current_user_id) try: while True: # Wait for a message OR a timeout (heartbeat) try: # Check for messages with a 20s timeout message = await asyncio.wait_for(queue.get(), timeout=20.0) yield f"data: {message}\n\n" except asyncio.TimeoutError: # Send a heartbeat comment to keep the connection alive yield ": heartbeat\n\n" finally: await notification_manager.unsubscribe(current_user_id, queue) return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", # Disable buffering for Nginx/Proxies }, )