engineHelper.py

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