import time import asyncio import threading from fastapi import WebSocket from datetime import datetime connections: dict[str, set[WebSocket]] = {} last_seen_ping: dict[str, float] = {} def parse_db_timestamp(ts_val) -> int: """Helper to convert various DB timestamp formats (Unix string or ISO string) to Unix integer.""" if not ts_val: return 0 if isinstance(ts_val, (int, float)): return int(ts_val) if isinstance(ts_val, str): if ts_val.isdigit(): return int(ts_val) try: # Handle ISO format from Supabase clean_ts = ts_val.replace('Z', '+00:00') dt = datetime.fromisoformat(clean_ts) return int(dt.timestamp()) except Exception: return 0 return 0 # ========================= # WS MANAGEMENT # ========================= async def register_ws(username: str, ws: WebSocket, connections_dict=None): if connections_dict is None: connections_dict = connections uname = username.lower() if uname not in connections_dict: connections_dict[uname] = set() connections_dict[uname].add(ws) def unregister_ws(username: str, ws: WebSocket, connections_dict=None): if connections_dict is None: connections_dict = connections uname = username.lower() if uname in connections_dict: connections_dict[uname].discard(ws) if not connections_dict[uname]: del connections_dict[uname] async def send_to_user(username: str, payload: dict, connections_dict=None): if connections_dict is None: connections_dict = connections uname = username.lower() if uname in connections_dict: for ws in list(connections_dict[uname]): try: await ws.send_json(payload) except Exception: connections_dict[uname].discard(ws) async def send_follow_notification(target_user, sender, action, connections_dict=None): if connections_dict is None: connections_dict = connections if target_user.lower() in connections_dict: if action == "follow": content = f"{sender} started following you" elif action == "unfollow": content = f"{sender} unfollowed you" elif action == "remove_follower": content = f"{sender} removed you as a follower" else: content = f"{sender} updated their connection with you" payload = { "type": "notification", "subtype": action, "sender": sender, "content": content, "timestamp": int(time.time()), } for ws in list(connections_dict[target_user.lower()]): try: await ws.send_json(payload) except Exception as e: connections_dict[target_user.lower()].discard(ws) # ========================= # BROADCAST # ========================= async def broadcast_like_update(post_id: int, likes: int, timeout: float = 2.0): message = {"type": "like_update", "post_id": post_id, "likes": likes} for user_sockets in list(connections.values()): for ws in list(user_sockets): asyncio.create_task(_safe_send(ws, message)) async def broadcast_online(): now = time.time() online_users = [ u for u, conns in connections.items() if conns ] payload = {"type": "online_list", "users": online_users} for user_sockets in list(connections.values()): for ws in list(user_sockets): asyncio.create_task(_safe_send(ws, payload)) async def broadcast_avatar_update(username: str, avatar: str): payload = {"type": "avatar_update", "username": username, "avatar": avatar} for user_sockets in list(connections.values()): for ws in list(user_sockets): asyncio.create_task(_safe_send(ws, payload)) async def broadcast_hero_update(username: str, hero: str): payload = {"type": "heroBanner_update", "username": username, "heroBanner": hero} for user_sockets in list(connections.values()): for ws in list(user_sockets): asyncio.create_task(_safe_send(ws, payload)) async def broadcast_theme_update(username: str, theme: int): payload = {"type": "theme_update", "username": username, "theme": theme} uname = username.lower() if uname in connections: for ws in list(connections[uname]): asyncio.create_task(_safe_send(ws, payload)) async def broadcast_to_all(payload: dict): """Broadcasts a payload to every active WebSocket connection in parallel.""" tasks = [] for user_sockets in list(connections.values()): for ws in list(user_sockets): tasks.append(asyncio.create_task(_safe_send(ws, payload))) if tasks: # We don't necessarily need to await all of them to finish if we want fire-and-forget, # but let's at least wait a short time for them to start. await asyncio.sleep(0) async def _safe_send(ws: WebSocket, payload: dict, timeout: float = 1.0): try: await asyncio.wait_for(ws.send_json(payload), timeout=timeout) except Exception: # Socket is likely dead, but we don't unregister here to avoid modification during iteration # The next ping/pong or cleanup cycle will handle it. pass async def evict_user_ws(username: str, connections_dict=None): """Sends a force-logout signal to all active connections of a user.""" if connections_dict is None: connections_dict = connections uname = username.lower() if uname in connections_dict: payload = {"type": "evict_user"} # Make a copy of the set to avoid modification during iteration for ws in list(connections_dict[uname]): try: await ws.send_json(payload) await ws.close(code=1000, reason="Account deleted") except Exception: pass # Clean up the connections dictionary if uname in connections_dict: del connections_dict[uname] if uname in last_seen_ping: del last_seen_ping[uname] def upload_media_to_cloudinary(data_url: str) -> str: """Upload a base64 data URL to Cloudinary, return the optimized public URL.""" import base64, uuid, mimetypes, cloudinary.uploader, os if not data_url or not data_url.startswith("data:"): return data_url try: import cloudinary cloudinary.config( cloud_name=os.environ.get("CLOUDINARY_CLOUD_NAME"), api_key=os.environ.get("CLOUDINARY_API_KEY"), api_secret=os.environ.get("CLOUDINARY_API_SECRET") ) header, b64 = data_url.split(",", 1) mime = header.split(";")[0].split(":")[1] ext = mimetypes.guess_extension(mime) or ".bin" name = f"{uuid.uuid4().hex}" # filename without extension for public_id file_bytes = base64.b64decode(b64) upload_res = cloudinary.uploader.upload( file_bytes, public_id=name, folder="yapstation_media", resource_type="auto" ) # Use 'f_auto,q_auto' for best performance/speed url = upload_res['secure_url'] optimized_url = url.replace("/upload/", "/upload/f_auto,q_auto/") return optimized_url except Exception as e: print(f"Cloudinary Upload Error: {e}") return data_url