WebSockets and Real-Time APIs

# HTTP: Client initiates, server responds, connection closes

# Request 1
Client → Server: GET /api/messages
Server → Client: 200 OK [message1, message2]
Connection: CLOSED

# Request 2 (new connection)
Client → Server: GET /api/messages
Server → Client: 200 OK [message1, message2, message3]
Connection: CLOSED

# Request 3 (new connection)
Client → Server: POST /api/messages {text: "Hello"}
Server → Client: 201 Created {id: 123}
Connection: CLOSED


# Limitations:
# - Server can't push data to client
# - Client must poll for updates
# - New TCP connection for each request (overhead)
# - Headers sent every time (Cookie, User-Agent, etc.)
# - Latency: connection setup + request + response

# WebSocket: Persistent connection, both sides can send anytime

# Handshake (HTTP upgrade)
Client → Server: GET /ws (Upgrade: websocket)
Server → Client: 101 Switching Protocols
Connection: ESTABLISHED (stays open)

# Bidirectional communication (no connection overhead)
Client → Server: {"action": "join", "room": "chat1"}
Server → Client: {"type": "joined", "room": "chat1"}

Server → Client: {"type": "message", "text": "New message!", "from": "Alice"}
Server → Client: {"type": "message", "text": "Another message!", "from": "Bob"}

Client → Server: {"action": "send", "text": "Hello everyone!"}
Server → Client: {"type": "message", "text": "Hello everyone!", "from": "You"}

# Connection stays open until explicitly closed
# No reconnection overhead
# Instant message delivery (no polling)
# Minimal bandwidth (no repeated headers)

# Step 1: Client initiates WebSocket handshake (HTTP request)

GET /ws HTTP/1.1
Host: api.example.com
Upgrade: websocket                          # Request protocol upgrade
Connection: Upgrade                         # Connection should be upgraded
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ== # Random base64 key
Sec-WebSocket-Version: 13                   # WebSocket protocol version
Origin: https://example.com                 # CORS origin


# Step 2: Server accepts upgrade (HTTP 101 response)

HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=  # Hashed key for validation


# Step 3: Connection upgraded to WebSocket
# - HTTP is done, WebSocket protocol takes over
# - Both sides can now send frames anytime
# - Connection stays open until explicitly closed


# WebSocket frame structure (simplified):
# ┌─────────┬──────────┬─────────┬──────────┐
# │ FIN=1   │ Opcode   │ Payload │ Data     │
# │ (final) │ (text/   │ Length  │ (actual  │
# │         │ binary)  │         │ message) │
# └─────────┴──────────┴─────────┴──────────┘

# Opcode types:
# 0x1 = Text frame
# 0x2 = Binary frame
# 0x8 = Close frame
# 0x9 = Ping frame
# 0xA = Pong frame (heartbeat)

# main.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """
    Simple echo server - sends back whatever it receives.

    Flow:
    1. Accept connection
    2. Receive message
    3. Echo it back
    4. Repeat until client disconnects
    """
    # Accept the WebSocket connection
    await websocket.accept()
    print("Client connected")

    try:
        while True:
            # Wait for message from client
            data = await websocket.receive_text()
            print(f"Received: {data}")

            # Echo message back to client
            await websocket.send_text(f"Echo: {data}")

    except WebSocketDisconnect:
        print("Client disconnected")

uvicorn main:app --reload

pip install websockets

# client.py
import asyncio
import websockets


async def test_echo():
    uri = "ws://localhost:8000/ws"
    async with websockets.connect(uri) as websocket:
        # Send message
        await websocket.send("Hello WebSocket!")

        # Receive echo
        response = await websocket.recv()
        print(f"Received: {response}")  # "Echo: Hello WebSocket!"


asyncio.run(test_echo())

# chat_server.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List
import json
from datetime import datetime, timezone

app = FastAPI()


class ConnectionManager:
    """Manage WebSocket connections and broadcasting"""

    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        """Accept and store new connection"""
        await websocket.accept()
        self.active_connections.append(websocket)
        print(f"Client connected. Total connections: {len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket):
        """Remove connection from active list"""
        self.active_connections.remove(websocket)
        print(f"Client disconnected. Total connections: {len(self.active_connections)}")

    async def broadcast(self, message: dict):
        """Send message to all connected clients"""
        for connection in self.active_connections:
            try:
                await connection.send_json(message)
            except Exception as e:
                print(f"Error sending to client: {e}")


manager = ConnectionManager()


@app.websocket("/ws/chat")
async def chat_endpoint(websocket: WebSocket):
    """
    Chat server supporting multiple clients.

    Features:
    - Broadcasts messages to all connected clients
    - Shows join/leave notifications
    - Includes timestamp and message ID
    """
    await manager.connect(websocket)

    # Notify all clients about new user
    await manager.broadcast({
        "type": "user_joined",
        "message": "A user joined the chat",
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "total_users": len(manager.active_connections)
    })

    try:
        while True:
            # Receive message from client
            data = await websocket.receive_text()

            # Parse message
            try:
                message_data = json.loads(data)
            except json.JSONDecodeError:
                # Plain text message
                message_data = {"text": data}

            # Broadcast to all clients
            broadcast_message = {
                "type": "message",
                "text": message_data.get("text", ""),
                "username": message_data.get("username", "Anonymous"),
                "timestamp": datetime.now(timezone.utc).isoformat()
            }

            await manager.broadcast(broadcast_message)

    except WebSocketDisconnect:
        manager.disconnect(websocket)

        # Notify all clients about user leaving
        await manager.broadcast({
            "type": "user_left",
            "message": "A user left the chat",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "total_users": len(manager.active_connections)
        })

uvicorn chat_server:app --reload

pip install websockets

# chat_client.py
import asyncio
import json
import websockets


async def chat_client(username: str):
    uri = "ws://localhost:8000/ws/chat"

    try:
        async with websockets.connect(uri) as websocket:

            # 1. Background task to receive messages
            async def receive_messages():
                try:
                    while True:
                        message = await websocket.recv()
                        data = json.loads(message)
                        # We use \r to clear the "You: " line momentarily
                        print(f"\r[{data['type']}] {data.get('text', data.get('message'))}\nYou: ", end="")
                except websockets.ConnectionClosed:
                    print("\nConnection closed by server.")

            asyncio.create_task(receive_messages())

            # 2. Non-blocking input loop
            while True:
                # Use run_in_executor to run the blocking input() in a thread
                text = await asyncio.get_event_loop().run_in_executor(
                    None, lambda: input("You: ")
                )

                if text.lower() == 'exit':
                    break

                message = json.dumps({"username": username, "text": text})
                await websocket.send(message)

    except Exception as e:
        print(f"Failed to connect: {e}")

# Terminal 1
python -c "import asyncio; from chat_client import chat_client; asyncio.run(chat_client('Alice'))"

# Terminal 2
python -c "import asyncio; from chat_client import chat_client; asyncio.run(chat_client('Bob'))"

# Terminal 3
python -c "import asyncio; from chat_client import chat_client; asyncio.run(chat_client('Charlie'))"

pip install fastapi pillow uvicorn

# image_server.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import io
from PIL import Image

app = FastAPI()


@app.websocket("/ws/images")
async def image_websocket(websocket: WebSocket):
    """
    WebSocket endpoint for binary image data.

    Receives images, processes them, sends back.
    """
    await websocket.accept()

    try:
        while True:
            # Receive binary data
            image_bytes = await websocket.receive_bytes()
            print(f"Received {len(image_bytes)} bytes")

            # Process image (example: resize)
            image = Image.open(io.BytesIO(image_bytes))
            resized = image.resize((100, 100))

            # Convert back to bytes
            output = io.BytesIO()
            resized.save(output, format='PNG')
            processed_bytes = output.getvalue()

            # Send binary data back
            await websocket.send_bytes(processed_bytes)
            print(f"Sent {len(processed_bytes)} bytes")

    except WebSocketDisconnect:
        print("Client disconnected")

uvicorn image_server:app --reload

pip install websockets

# image_client.py
import asyncio
import websockets


async def send_image():
    uri = "ws://localhost:8000/ws/images"

    async with websockets.connect(uri) as websocket:
        # Read image file
        with open("photo.jpg", "rb") as f:
            image_data = f.read()

        # Send as binary
        await websocket.send(image_data)

        # Receive processed image
        processed = await websocket.recv()

        # Save result
        with open("photo_resized.png", "wb") as f:
            f.write(processed)

        print("Image processed and saved!")


asyncio.run(send_image())

pip install fastapi pyjwt uvicorn

# auth_server.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, status
import jwt
from datetime import datetime, timedelta, timezone

app = FastAPI()

# Secret key for JWT
SECRET_KEY = "your-secret-key"
ALGORITHM = "HS256"


def create_access_token(user_id: str) -> str:
    """Generate JWT token"""
    payload = {
        "user_id": user_id,
        "exp": datetime.now(timezone.utc) + timedelta(hours=1)
    }
    return jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)


def verify_token(token: str) -> dict:
    """Verify JWT token and return payload"""
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return payload
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Token expired")
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="Invalid token")


@app.websocket("/ws/chat")
async def authenticated_chat(websocket: WebSocket, token: str):
    """
    WebSocket endpoint with authentication.

    Usage:
    ws://localhost:8000/ws/chat?token=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...

    The token parameter is extracted from query string.
    """
    # Verify token before accepting connection
    try:
        payload = verify_token(token)
        user_id = payload["user_id"]
    except HTTPException as e:
        # Reject connection with close code
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason=e.detail)
        return

    # Token valid - accept connection
    await websocket.accept()
    print(f"Authenticated user: {user_id}")

    try:
        while True:
            data = await websocket.receive_text()

            # Include user_id with message
            response = {
                "user_id": user_id,
                "message": data,
                "timestamp": datetime.now(timezone.utc).isoformat()
            }

            await websocket.send_json(response)

    except WebSocketDisconnect:
        print(f"User {user_id} disconnected")

uvicorn auth_server:app --reload

pip install websockets requests

# auth_client.py
import asyncio
import websockets
import requests


async def connect_with_auth():
    # Get token from login endpoint
    # Note: You'll need to implement a /login endpoint that returns a token
    response = requests.post("http://localhost:8000/login", json={
        "username": "alice",
        "password": "secret"
    })
    token = response.json()["access_token"]

    # Connect WebSocket with token in query parameter
    uri = f"ws://localhost:8000/ws/chat?token={token}"

    async with websockets.connect(uri) as websocket:
        await websocket.send("Hello!")
        response = await websocket.recv()
        print(response)  # Includes authenticated user_id


asyncio.run(connect_with_auth())

import redis
import time
from collections import defaultdict

redis_client = redis.Redis(host='localhost', port=6379, decode_responses=True)

# In-memory rate limiter (use Redis for production)
message_timestamps = defaultdict(list)


async def check_connection_limit(user_id: str, max_connections: int = 5) -> bool:
    """
    Check if user has exceeded concurrent connection limit.

    Args:
        user_id: User identifier
        max_connections: Maximum concurrent connections per user

    Returns:
        True if allowed, False if limit exceeded
    """
    key = f"ws_connections:{user_id}"
    connection_count = redis_client.get(key)

    if connection_count and int(connection_count) >= max_connections:
        return False

    return True


async def check_message_rate_limit(user_id: str, limit: int = 10, window: int = 60) -> bool:
    """
    Rate limit messages per user.

    Args:
        user_id: User identifier
        limit: Max messages per window
        window: Time window in seconds

    Returns:
        True if allowed, False if rate limited
    """
    current_time = time.time()
    window_start = current_time - window

    # Get user's message timestamps
    timestamps = message_timestamps[user_id]

    # Remove old timestamps
    timestamps[:] = [ts for ts in timestamps if ts > window_start]

    # Check limit
    if len(timestamps) >= limit:
        return False

    # Add current timestamp
    timestamps.append(current_time)
    return True


@app.websocket("/ws/chat")
async def rate_limited_chat(websocket: WebSocket, token: str):
    """Chat endpoint with rate limiting"""

    # Authenticate
    try:
        payload = verify_token(token)
        user_id = payload["user_id"]
    except HTTPException:
        await websocket.close(code=1008, reason="Authentication failed")
        return

    # Check connection limit
    if not await check_connection_limit(user_id, max_connections=5):
        await websocket.close(code=1008, reason="Too many connections")
        return

    # Increment connection count
    connection_key = f"ws_connections:{user_id}"
    redis_client.incr(connection_key)
    redis_client.expire(connection_key, 3600)

    await websocket.accept()

    try:
        while True:
            data = await websocket.receive_text()

            # Rate limit messages
            if not await check_message_rate_limit(user_id, limit=10, window=60):
                await websocket.send_json({
                    "type": "error",
                    "message": "Rate limit exceeded. Max 10 messages per minute."
                })
                continue

            # Process message
            await websocket.send_json({
                "type": "message",
                "text": data,
                "user_id": user_id
            })

    except WebSocketDisconnect:
        # Decrement connection count
        redis_client.decr(connection_key)


# Result:
# User can open max 5 WebSocket connections
# User can send max 10 messages per minute
# Exceeded limits get error message or connection closed

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import Dict, List, Set
from collections import defaultdict
import json

app = FastAPI()


class RoomManager:
    """
    Manage WebSocket rooms for group chat.

    Features:
    - Join/leave rooms
    - Broadcast to specific room
    - List users in room
    - Private messages
    """

    def __init__(self):
        # room_id -> set of WebSocket connections
        self.rooms: Dict[str, Set[WebSocket]] = defaultdict(set)
        # websocket -> user_info
        self.users: Dict[WebSocket, dict] = {}

    async def join_room(self, websocket: WebSocket, room_id: str, username: str):
        """Add user to room"""
        self.rooms[room_id].add(websocket)
        self.users[websocket] = {
            "username": username,
            "room": room_id
        }

        # Notify room
        await self.broadcast_to_room(room_id, {
            "type": "user_joined",
            "username": username,
            "room": room_id,
            "users_count": len(self.rooms[room_id])
        }, exclude=websocket)

    async def leave_room(self, websocket: WebSocket):
        """Remove user from room"""
        if websocket not in self.users:
            return

        user_info = self.users[websocket]
        room_id = user_info["room"]
        username = user_info["username"]

        # Remove from room
        self.rooms[room_id].discard(websocket)
        del self.users[websocket]

        # Notify room
        await self.broadcast_to_room(room_id, {
            "type": "user_left",
            "username": username,
            "room": room_id,
            "users_count": len(self.rooms[room_id])
        })

    async def broadcast_to_room(self, room_id: str, message: dict, exclude: WebSocket = None):
        """Send message to all users in room"""
        if room_id not in self.rooms:
            return

        disconnected = []

        for connection in self.rooms[room_id]:
            if connection == exclude:
                continue

            try:
                await connection.send_json(message)
            except Exception:
                disconnected.append(connection)

        # Clean up disconnected clients
        for conn in disconnected:
            await self.leave_room(conn)

    def get_room_users(self, room_id: str) -> List[str]:
        """Get list of usernames in room"""
        if room_id not in self.rooms:
            return []

        return [
            self.users[ws]["username"]
            for ws in self.rooms[room_id]
            if ws in self.users
        ]


manager = RoomManager()


@app.websocket("/ws/rooms")
async def room_websocket(websocket: WebSocket, username: str):
    """
    Room-based chat WebSocket.

    Messages format:
    - Join: {"action": "join", "room": "general"}
    - Send: {"action": "message", "text": "Hello!"}
    - Leave: {"action": "leave"}
    """
    await websocket.accept()
    current_room = None

    try:
        while True:
            # Receive message
            data = await websocket.receive_text()
            message = json.loads(data)

            action = message.get("action")

            if action == "join":
                room_id = message.get("room", "general")

                # Leave current room if in one
                if current_room:
                    await manager.leave_room(websocket)

                # Join new room
                await manager.join_room(websocket, room_id, username)
                current_room = room_id

                # Send confirmation
                await websocket.send_json({
                    "type": "joined",
                    "room": room_id,
                    "users": manager.get_room_users(room_id)
                })

            elif action == "message":
                if not current_room:
                    await websocket.send_json({
                        "type": "error",
                        "message": "Not in a room. Join a room first."
                    })
                    continue

                # Broadcast message to room
                await manager.broadcast_to_room(current_room, {
                    "type": "message",
                    "username": username,
                    "text": message.get("text", ""),
                    "room": current_room
                })

            elif action == "leave":
                if current_room:
                    await manager.leave_room(websocket)
                    current_room = None

                    await websocket.send_json({
                        "type": "left",
                        "message": "Left the room"
                    })

            elif action == "list_users":
                if current_room:
                    users = manager.get_room_users(current_room)
                    await websocket.send_json({
                        "type": "user_list",
                        "users": users,
                        "count": len(users)
                    })

    except WebSocketDisconnect:
        if current_room:
            await manager.leave_room(websocket)


# Client example
async def chat_in_rooms():
    uri = "ws://localhost:8000/ws/rooms?username=Alice"

    async with websockets.connect(uri) as websocket:
        # Join room
        await websocket.send(json.dumps({
            "action": "join",
            "room": "python-devs"
        }))

        # Receive confirmation
        response = json.loads(await websocket.recv())
        print(f"Joined {response['room']}, users: {response['users']}")

        # Send message to room
        await websocket.send(json.dumps({
            "action": "message",
            "text": "Hello Python developers!"
        }))

        # Switch room
        await websocket.send(json.dumps({
            "action": "join",
            "room": "javascript-devs"
        }))

        # Send to new room
        await websocket.send(json.dumps({
            "action": "message",
            "text": "Hello JS developers!"
        }))

import asyncio
from datetime import datetime, timezone


class PresenceManager:
    """Track online/offline status and last seen"""

    def __init__(self):
        self.online_users: Dict[str, dict] = {}

    async def set_online(self, user_id: str, websocket: WebSocket):
        """Mark user as online"""
        self.online_users[user_id] = {
            "status": "online",
            "websocket": websocket,
            "connected_at": datetime.now(timezone.utc).isoformat(),
            "last_activity": datetime.now(timezone.utc).isoformat()
        }

        # Broadcast presence update
        await self.broadcast_presence_update(user_id, "online")

    async def set_offline(self, user_id: str):
        """Mark user as offline"""
        if user_id in self.online_users:
            del self.online_users[user_id]

            # Broadcast presence update
            await self.broadcast_presence_update(user_id, "offline")

    async def update_activity(self, user_id: str):
        """Update last activity timestamp"""
        if user_id in self.online_users:
            self.online_users[user_id]["last_activity"] = datetime.now(timezone.utc).isoformat()

    def get_online_users(self) -> List[dict]:
        """Get list of online users"""
        return [
            {
                "user_id": user_id,
                "status": info["status"],
                "connected_at": info["connected_at"]
            }
            for user_id, info in self.online_users.items()
        ]

    async def broadcast_presence_update(self, user_id: str, status: str):
        """Notify all users about presence change"""
        message = {
            "type": "presence",
            "user_id": user_id,
            "status": status,
            "timestamp": datetime.now(timezone.utc).isoformat()
        }

        for user_info in self.online_users.values():
            try:
                await user_info["websocket"].send_json(message)
            except Exception:
                pass


presence = PresenceManager()


@app.websocket("/ws/presence")
async def presence_websocket(websocket: WebSocket, user_id: str):
    """WebSocket with presence tracking"""
    await websocket.accept()

    # Mark user as online
    await presence.set_online(user_id, websocket)

    # Send current online users
    await websocket.send_json({
        "type": "online_users",
        "users": presence.get_online_users()
    })

    try:
        while True:
            data = await websocket.receive_text()

            # Update last activity
            await presence.update_activity(user_id)

            # Handle message...

    except WebSocketDisconnect:
        # Mark user as offline
        await presence.set_offline(user_id)


# All connected clients receive:
# {"type": "presence", "user_id": "alice", "status": "online"}
# {"type": "presence", "user_id": "bob", "status": "online"}
# {"type": "presence", "user_id": "alice", "status": "offline"}

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio
import json
import time

app = FastAPI()


async def event_generator():
    """
    Generate server-sent events.

    SSE format:
    data: message



    With event type:
    event: notification

    data: message


    """
    counter = 0

    while True:
        # Simulate new data
        counter += 1

        # Format as SSE event
        data = json.dumps({
            "counter": counter,
            "timestamp": time.time(),
            "message": f"Update #{counter}"
        })

        # SSE format: "data: " prefix, double newline at end
        yield f"data: {data}

"

        # Send every 2 seconds
        await asyncio.sleep(2)


@app.get("/sse/events")
async def sse_endpoint():
    """
    Server-Sent Events endpoint.

    Returns a stream of events to the client.
    Client automatically reconnects if disconnected.
    """
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"  # Disable nginx buffering
        }
    )


# JavaScript client:
# const eventSource = new EventSource('http://localhost:8000/sse/events');
#
# eventSource.onmessage = (event) => {
#     const data = JSON.parse(event.data);
#     console.log('Received:', data);
# };
#
# eventSource.onerror = (error) => {
#     console.error('SSE error:', error);
#     // EventSource auto-reconnects after 3 seconds
# };


# Python client
import requests

def sse_client():
    """Receive server-sent events"""
    response = requests.get(
        'http://localhost:8000/sse/events',
        stream=True
    )

    for line in response.iter_lines():
        if line:
            line = line.decode('utf-8')
            if line.startswith('data: '):
                data = json.loads(line[6:])  # Remove "data: " prefix
                print(f"Received: {data}")

async def notification_generator():
    """
    SSE with event types and IDs.

    Event types allow different handlers:
    - eventSource.addEventListener('notification', handler)
    - eventSource.addEventListener('alert', handler)

    Event IDs enable resuming from last received event.
    """
    event_id = 0

    while True:
        event_id += 1

        # Randomly choose event type
        if event_id % 3 == 0:
            event_type = "alert"
            data = {"severity": "high", "message": "Important alert!"}
        else:
            event_type = "notification"
            data = {"message": f"Notification #{event_id}"}

        # SSE format with event type and ID
        yield f"id: {event_id}
"  # Event ID for resume
        yield f"event: {event_type}
"  # Event type
        yield f"data: {json.dumps(data)}

"  # Data (double newline at end)

        await asyncio.sleep(2)


@app.get("/sse/notifications")
async def sse_notifications(last_event_id: str = None):
    """
    SSE endpoint with resume support.

    If client reconnects, they send Last-Event-ID header.
    Server can resume from that point.
    """
    # Use last_event_id to resume from specific event
    # (implementation depends on your event storage)

    return StreamingResponse(
        notification_generator(),
        media_type="text/event-stream"
    )


# JavaScript client with event types:
# const eventSource = new EventSource('/sse/notifications');
#
# eventSource.addEventListener('notification', (event) => {
#     const data = JSON.parse(event.data);
#     console.log('Notification:', data);
# });
#
# eventSource.addEventListener('alert', (event) => {
#     const data = JSON.parse(event.data);
#     console.warn('ALERT:', data);
# });
#
# // On reconnect, browser sends Last-Event-ID header automatically
# // Server can resume from that event

from fastapi import FastAPI, WebSocket
import redis.asyncio as redis
import json
import asyncio
import time

app = FastAPI()

# Redis connection for pub/sub
redis_client = None


class ScalableConnectionManager:
    """
    WebSocket manager with Redis pub/sub for scaling.

    Architecture:
    - Each server has local WebSocket connections
    - Messages published to Redis channel
    - All servers subscribe to Redis
    - Each server broadcasts to local connections
    """

    def __init__(self):
        self.active_connections: List[WebSocket] = []
        self.redis_client = None
        self.pubsub = None

    async def connect(self, websocket: WebSocket):
        """Accept WebSocket connection"""
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        """Remove WebSocket connection"""
        self.active_connections.remove(websocket)

    async def setup_redis(self):
        """Setup Redis connection and pub/sub"""
        self.redis_client = await redis.from_url("redis://localhost:6379")
        self.pubsub = self.redis_client.pubsub()

        # Subscribe to broadcast channel
        await self.pubsub.subscribe("chat_messages")

        # Start listening in background
        asyncio.create_task(self.listen_to_redis())

    async def listen_to_redis(self):
        """
        Listen to Redis pub/sub messages.

        When message received from Redis:
        1. Parse message
        2. Broadcast to all local WebSocket connections
        """
        async for message in self.pubsub.listen():
            if message["type"] == "message":
                # Parse message
                data = json.loads(message["data"])

                # Broadcast to all local connections
                await self.broadcast_local(data)

    async def broadcast_local(self, message: dict):
        """Broadcast to local WebSocket connections only"""
        disconnected = []

        for connection in self.active_connections:
            try:
                await connection.send_json(message)
            except Exception:
                disconnected.append(connection)

        # Clean up
        for conn in disconnected:
            self.disconnect(conn)

    async def broadcast_global(self, message: dict):
        """
        Broadcast to ALL servers via Redis.

        Steps:
        1. Publish message to Redis channel
        2. Redis delivers to all subscribed servers
        3. Each server broadcasts to local connections
        """
        # Publish to Redis
        await self.redis_client.publish(
            "chat_messages",
            json.dumps(message)
        )


manager = ScalableConnectionManager()


@app.on_event("startup")
async def startup():
    """Setup Redis on app startup"""
    await manager.setup_redis()


@app.on_event("shutdown")
async def shutdown():
    """Cleanup Redis on shutdown"""
    if manager.redis_client:
        await manager.redis_client.close()


@app.websocket("/ws/chat")
async def scalable_chat(websocket: WebSocket):
    """
    Scalable chat endpoint.

    Works across multiple servers via Redis pub/sub.
    """
    await manager.connect(websocket)

    try:
        while True:
            # Receive message from client
            data = await websocket.receive_text()

            # Create message
            message = {
                "type": "message",
                "text": data,
                "timestamp": time.time()
            }

            # Broadcast to ALL servers (via Redis)
            await manager.broadcast_global(message)

    except WebSocketDisconnect:
        manager.disconnect(websocket)


# Flow with 2 servers:
# 1. Alice (Server A) sends "Hello"
# 2. Server A publishes to Redis: {"text": "Hello"}
# 3. Redis notifies Server A and Server B
# 4. Server A broadcasts to local connections (Alice)
# 5. Server B broadcasts to local connections (Bob)
# 6. Both Alice and Bob receive "Hello"!

# Nginx configuration for WebSocket load balancing

upstream websocket_servers {
    # Use IP hash for sticky sessions
    # Same client always routed to same server
    ip_hash;

    server 10.0.1.1:8000;
    server 10.0.1.2:8000;
    server 10.0.1.3:8000;
}

server {
    listen 80;
    server_name api.example.com;

    location /ws/ {
        # WebSocket proxy settings
        proxy_pass http://websocket_servers;
        proxy_http_version 1.1;

        # Required for WebSocket upgrade
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection "upgrade";

        # Forward headers
        proxy_set_header Host $host;
        proxy_set_header X-Real-IP $remote_addr;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;

        # Timeouts
        proxy_read_timeout 86400s;  # 24 hours
        proxy_send_timeout 86400s;

        # Disable buffering
        proxy_buffering off;
    }
}


# Why ip_hash?
# - Same client connects to same server
# - Reduces Redis traffic (local broadcasts more common)
# - Better for presence tracking (user's connections on one server)
#
# Alternative: least_conn (connect to least busy server)
# Better distribution but more Redis traffic

import asyncio
from fastapi import WebSocket

HEARTBEAT_INTERVAL = 30  # Send ping every 30 seconds
HEARTBEAT_TIMEOUT = 60   # Close if no pong in 60 seconds


async def heartbeat_handler(websocket: WebSocket):
    """
    Send periodic ping frames.

    If client doesn't respond with pong within timeout, close connection.
    """
    while True:
        try:
            # Send ping
            await websocket.send_text('{"type": "ping"}')

            # Wait for pong (with timeout)
            try:
                pong = await asyncio.wait_for(
                    websocket.receive_text(),
                    timeout=HEARTBEAT_TIMEOUT
                )

                # Check if it's a pong response
                if '"type": "pong"' not in pong:
                    # Not a pong, put message back in queue
                    pass

            except asyncio.TimeoutError:
                # No pong received, connection is dead
                print("Heartbeat timeout - closing connection")
                await websocket.close(code=1000, reason="Heartbeat timeout")
                break

            # Wait before next ping
            await asyncio.sleep(HEARTBEAT_INTERVAL)

        except Exception as e:
            print(f"Heartbeat error: {e}")
            break


@app.websocket("/ws")
async def websocket_with_heartbeat(websocket: WebSocket):
    """WebSocket with heartbeat monitoring"""
    await websocket.accept()

    # Start heartbeat in background
    heartbeat_task = asyncio.create_task(heartbeat_handler(websocket))

    try:
        while True:
            data = await websocket.receive_text()

            # Handle pong responses
            message = json.loads(data)
            if message.get("type") == "pong":
                continue  # Just acknowledgment, no processing

            # Process regular message
            await websocket.send_text(f"Echo: {data}")

    except WebSocketDisconnect:
        heartbeat_task.cancel()
        print("Client disconnected")


# Client should respond to pings:
# websocket.onmessage = (event) => {
#     const data = JSON.parse(event.data);
#     if (data.type === 'ping') {
#         websocket.send(JSON.stringify({type: 'pong'}));
#     }
# };

from collections import defaultdict

# Global limits
MAX_CONNECTIONS_PER_USER = 5
MAX_TOTAL_CONNECTIONS = 10000
MAX_MESSAGE_SIZE = 64 * 1024  # 64 KB

# Track connections
connection_counts = defaultdict(int)
total_connections = 0


async def check_connection_limits(user_id: str) -> bool:
    """Check if user can open new connection"""
    global total_connections

    # Check total server limit
    if total_connections >= MAX_TOTAL_CONNECTIONS:
        return False

    # Check per-user limit
    if connection_counts[user_id] >= MAX_CONNECTIONS_PER_USER:
        return False

    return True


@app.websocket("/ws")
async def limited_websocket(websocket: WebSocket, user_id: str):
    """WebSocket with connection limits"""
    global total_connections

    # Check limits
    if not await check_connection_limits(user_id):
        await websocket.close(
            code=1008,
            reason="Connection limit exceeded"
        )
        return

    # Accept connection
    await websocket.accept()

    # Track connection
    connection_counts[user_id] += 1
    total_connections += 1

    try:
        while True:
            # Receive with size limit
            data = await websocket.receive_text()

            if len(data) > MAX_MESSAGE_SIZE:
                await websocket.send_json({
                    "type": "error",
                    "message": "Message too large (max 64KB)"
                })
                continue

            # Process message
            await websocket.send_text(f"Echo: {data}")

    except WebSocketDisconnect:
        pass
    finally:
        # Cleanup
        connection_counts[user_id] -= 1
        total_connections -= 1
        print(f"Total connections: {total_connections}")


# Monitoring endpoint
@app.get("/metrics/websockets")
async def websocket_metrics():
    """Get WebSocket connection metrics"""
    return {
        "total_connections": total_connections,
        "max_connections": MAX_TOTAL_CONNECTIONS,
        "usage_percent": (total_connections / MAX_TOTAL_CONNECTIONS) * 100,
        "connections_by_user": dict(connection_counts)
    }

import signal

# Track active connections for graceful shutdown
active_websockets: Set[WebSocket] = set()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket with graceful shutdown support"""
    await websocket.accept()
    active_websockets.add(websocket)

    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Echo: {data}")

    except WebSocketDisconnect:
        pass
    finally:
        active_websockets.discard(websocket)


async def graceful_shutdown():
    """
    Gracefully close all WebSocket connections.

    Notify clients that server is shutting down,
    give them time to reconnect to another server.
    """
    print(f"Shutting down {len(active_websockets)} WebSocket connections...")

    # Notify all clients
    for ws in active_websockets:
        try:
            await ws.send_json({
                "type": "server_shutdown",
                "message": "Server restarting, please reconnect",
                "reconnect_delay": 5  # Wait 5 seconds before reconnecting
            })

            # Close gracefully
            await ws.close(code=1001, reason="Server going away")

        except Exception as e:
            print(f"Error closing WebSocket: {e}")

    print("All WebSocket connections closed")


@app.on_event("shutdown")
async def shutdown_event():
    """Handle application shutdown"""
    await graceful_shutdown()


# For deployment with zero downtime:
# 1. Deploy new server instances
# 2. Update load balancer to route new connections to new servers
# 3. Old servers finish handling existing connections
# 4. Gracefully shut down old servers after 5-10 minutes