145 lines
4.5 KiB
Python
145 lines
4.5 KiB
Python
import json
|
|
import logging
|
|
from datetime import datetime, timezone, timedelta
|
|
|
|
from fastapi import APIRouter, Depends, Query, WebSocket, WebSocketDisconnect
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.dependencies import get_db, get_current_user
|
|
from app.models.user import User
|
|
from app.models.notification import Notification
|
|
from app.schemas.notification import NotificationListOut, NotificationCreditsOut, UnreadCountOut
|
|
from app.services.auth import decode_access_token
|
|
from app.services.notification import (
|
|
get_notifications,
|
|
mark_read,
|
|
mark_all_read,
|
|
get_unread_count,
|
|
create_notification,
|
|
)
|
|
|
|
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()
|
|
|
|
router = APIRouter(prefix="/notifications", tags=["notifications"])
|
|
|
|
# Active WebSocket connections: user_id -> set of WebSockets
|
|
_active_connections: dict[str, set[WebSocket]] = {}
|
|
|
|
|
|
async def push_notification_to_user(user_id: str, notification: Notification) -> None:
|
|
"""Push a notification to all active WebSocket connections for a user."""
|
|
connections = _active_connections.get(user_id, set())
|
|
if not connections:
|
|
return
|
|
message = json.dumps({
|
|
"id": notification.id,
|
|
"title": notification.title,
|
|
"content": notification.content,
|
|
"type": notification.type,
|
|
"is_read": notification.is_read,
|
|
"created_at": _to_local_str(notification.created_at),
|
|
}, ensure_ascii=False)
|
|
dead: list[WebSocket] = []
|
|
for ws in connections:
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception:
|
|
dead.append(ws)
|
|
for ws in dead:
|
|
connections.discard(ws)
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def notifications_ws(websocket: WebSocket):
|
|
"""Authenticated WebSocket for real-time notifications."""
|
|
await websocket.accept()
|
|
|
|
# Authenticate via token query param or first message
|
|
token = websocket.query_params.get("token")
|
|
if not token:
|
|
try:
|
|
data = await websocket.receive_text()
|
|
msg = json.loads(data)
|
|
token = msg.get("token")
|
|
except Exception:
|
|
await websocket.close(code=4001, reason="Authentication required")
|
|
return
|
|
|
|
user_id = decode_access_token(token) if token else None
|
|
if not user_id or user_id.startswith("captcha:"):
|
|
await websocket.close(code=4001, reason="Invalid token")
|
|
return
|
|
|
|
# Register connection
|
|
if user_id not in _active_connections:
|
|
_active_connections[user_id] = set()
|
|
_active_connections[user_id].add(websocket)
|
|
|
|
try:
|
|
# Keep connection alive, listen for pings
|
|
while True:
|
|
data = await websocket.receive_text()
|
|
# Echo back for heartbeat
|
|
if data == "ping":
|
|
await websocket.send_text("pong")
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
conns = _active_connections.get(user_id)
|
|
if conns:
|
|
conns.discard(websocket)
|
|
if not conns:
|
|
_active_connections.pop(user_id, None)
|
|
|
|
|
|
@router.get("", response_model=NotificationListOut)
|
|
async def list_notifications(
|
|
page: int = Query(1, ge=1),
|
|
page_size: int = Query(20, ge=1, le=100),
|
|
is_read: bool = Query(None),
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
items, total = await get_notifications(db, current_user.id, page, page_size, is_read)
|
|
credits = NotificationCreditsOut(balance=round(float(current_user.credits or 0.0), 2))
|
|
return {"items": items, "total": total, "credits": credits}
|
|
|
|
|
|
@router.put("/{notification_id}/read")
|
|
async def read_notification(
|
|
notification_id: str,
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
await mark_read(db, notification_id, current_user.id)
|
|
return {"message": "ok"}
|
|
|
|
|
|
@router.put("/read-all")
|
|
async def read_all(
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
await mark_all_read(db, current_user.id)
|
|
return {"message": "ok"}
|
|
|
|
|
|
@router.get("/unread-count", response_model=UnreadCountOut)
|
|
async def unread_count(
|
|
current_user: User = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
count = await get_unread_count(db, current_user.id)
|
|
return UnreadCountOut(count=count)
|