server.py

'''
asyncio module let us handle asynchronous tasks. 
connected_clients set keeps track of all connected clients. 
Each time a new client joins, the handle_client function adds them to that set, listens for anything they send and then passes their messages along to everyone else. 
Main function sets up the localhost at a port ... and waits for people to connect
Call asyncio.run(main()) to start the server when the script is executed directly.
'''
import asyncio
import websockets
import jwt
from datetime import datetime, timedelta
from sqlalchemy import create_engine, Integer, String, Column, DateTime
from sqlalchemy.orm import declarative_base, sessionmaker

# ---------------------------
# Config / DB setup
# ---------------------------
SECRET_KEY = "supersecretkey"
engine = create_engine("sqlite:///messages.db", echo=True)
Base = declarative_base()
SessionLocal = sessionmaker(bind=engine)

# ---------------------------
# Database model
# ---------------------------
class Message(Base):
    __tablename__ = "messages"
    id = Column(Integer, primary_key=True)
    username = Column(String, nullable=False)
    content = Column(String, nullable=False)
    timestamp = Column(DateTime, default=datetime.now)

# create tables
Base.metadata.create_all(engine)

connected_clients = {}

# ---------------------------
# JWT helper functions
# ---------------------------

def verify_token(token):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
        return payload["username"]
    except:
        return None

# ---------------------------
# WebSocket handler
# ---------------------------
async def handle_client(websocket):
    # Extract JWT from query string: ws://localhost:5000/?token=...
    path = websocket.request.path
    token = None
    if "?" in path:
        query = path.split("?")[1]
        for kv in query.split("&"):
            if "=" in kv:
                k, v = kv.split("=")
                if k == "token":
                    token = v
                    break

    username = verify_token(token)
    if not username:
        await websocket.close()
        return

    # Add websocket -> username mapping
    connected_clients[websocket] = username

    # Send last 10 messages from DB to the client
    db = SessionLocal()
    messages = db.query(Message).order_by(Message.timestamp.desc()).limit(10).all()
    for msg in reversed(messages):
        await websocket.send(f"{msg.username}: {msg.content}")

    try:
        async for message_text in websocket:
            try:
                # Save message to DB
                new_message = Message(username=username, content=message_text)
                db.add(new_message)
                db.commit()
            except Exception as e:
                db.rollback()
                print("DB error:",e)

            # Broadcast to all connected clients
            send_tasks = []
            for client in list(connected_clients):
                send_tasks.append(client.send(f"{username}: {message_text}"))
            await asyncio.gather(*send_tasks)
    finally:
        # Remove client
        connected_clients.pop(websocket, None)
        db.close()

# ---------------------------
# Main server function
# ---------------------------
async def main():
    async with websockets.serve(handle_client,"localhost",5000,ping_interval=30,ping_timeout=30):
        print("Server started on ws://localhost:5000")
        await asyncio.Future()  # run forever

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("Server is shutting down")
    finally:
        print("Clean up complete")