Design a Chat System

import asyncio
import websockets
import json
from collections import defaultdict

# Map: user_id -> set of websocket connections (multi-device)
connections: dict[str, set] = defaultdict(set)

async def handle_connection(websocket, path):
    """Handle a new WebSocket connection."""
    user_id = await authenticate(websocket)
    connections[user_id].add(websocket)

    try:
        async for raw_message in websocket:
            message = json.loads(raw_message)
            await process_message(user_id, message)
    finally:
        connections[user_id].discard(websocket)
        if not connections[user_id]:
            del connections[user_id]
            await update_presence(user_id, "offline")


async def process_message(sender_id: str, message: dict):
    """Route message to recipient(s)."""
    msg_type = message["type"]

    if msg_type == "direct":
        recipient_id = message["to"]
        payload = {
            "from": sender_id,
            "text": message["text"],
            "timestamp": get_timestamp(),
            "msg_id": generate_msg_id()
        }

        # Persist to message store
        await store_message(sender_id, recipient_id, payload)

        # Deliver to online recipient (all devices)
        if recipient_id in connections:
            for ws in connections[recipient_id]:
                await ws.send(json.dumps(payload))
        else:
            # User offline: send push notification
            await send_push_notification(recipient_id, payload)

    elif msg_type == "group":
        await handle_group_message(sender_id, message)


# Start WebSocket server
server = websockets.serve(handle_connection, "0.0.0.0", 8765)
asyncio.get_event_loop().run_until_complete(server)
asyncio.get_event_loop().run_forever()

import redis
import time

class PresenceService:
    """
    Tracks online/offline status using Redis with TTL.

    Approach: Heartbeat-based presence detection.
    - Client sends heartbeat every 5 seconds
    - If no heartbeat for 30 seconds, user is offline
    - Uses Redis key expiry for automatic cleanup
    """

    def __init__(self):
        self.redis = redis.Redis(host='localhost', port=6379)
        self.HEARTBEAT_TTL = 30  # seconds

    def heartbeat(self, user_id: str):
        """Update user's last-seen timestamp with TTL."""
        key = f"presence:{user_id}"
        self.redis.setex(key, self.HEARTBEAT_TTL, str(time.time()))

    def is_online(self, user_id: str) -> bool:
        """Check if user has an active heartbeat."""
        return self.redis.exists(f"presence:{user_id}") > 0

    def get_last_seen(self, user_id: str) -> float:
        """Get user's last heartbeat timestamp."""
        ts = self.redis.get(f"presence:{user_id}")
        return float(ts) if ts else 0

    def get_online_friends(self, user_id: str) -> list:
        """Get list of online friends using pipeline for efficiency."""
        friends = get_friends_list(user_id)
        pipe = self.redis.pipeline()
        for friend_id in friends:
            pipe.exists(f"presence:{friend_id}")
        results = pipe.execute()
        return [f for f, online in zip(friends, results) if online]

-- Cassandra schema for chat messages
-- Partition key: chat_id (groups messages by conversation)
-- Clustering key: msg_id (sorts messages within conversation)

CREATE TABLE messages (
    chat_id     UUID,
    msg_id      TIMEUUID,
    sender_id   UUID,
    content     TEXT,
    msg_type    TEXT,      -- 'text', 'image', 'file'
    created_at  TIMESTAMP,
    PRIMARY KEY (chat_id, msg_id)
) WITH CLUSTERING ORDER BY (msg_id DESC);

-- Efficient query: "Get latest 50 messages in chat"
-- SELECT * FROM messages WHERE chat_id = ? LIMIT 50;

-- User's chat list (most recent message per chat)
CREATE TABLE user_chats (
    user_id     UUID,
    chat_id     UUID,
    last_msg    TEXT,
    last_time   TIMESTAMP,
    unread      INT,
    PRIMARY KEY (user_id, last_time)
) WITH CLUSTERING ORDER BY (last_time DESC);