app.py

from dotenv import load_dotenv
from flask import Flask, redirect, render_template, request, session, url_for
from flask_caching import Cache
from flask_socketio import join_room, leave_room, SocketIO, send
from string import ascii_letters
import os
import random


dotenv_path = os.path.join(os.path.dirname(__file__), ".env")
load_dotenv(dotenv_path)

ROOM_CODE_LENGTH = 8
AN_HOUR_IN_SEC = 3600

app = Flask(__name__)
app.config["SECRET_KEY"] = os.environ.get("SECRET_KEY")
cache = Cache(config={
    "CACHE_TYPE": "redis",
    "CACHE_REDIS_URL": os.environ.get("CACHE_REDIS_URL"),
})

socketio = SocketIO(app)
cache.init_app(app)


@app.route("/", methods=["GET", "POST"])
def home():
    session.clear()
    if request.method == "POST":
        username = request.form.get("username")
        room_code = request.form.get("room_code")

        if not username or not room_code:
            return render_template("home.html", error="Userame and Room Code are required")

        if not cache.get(room_code):
            return render_template("home.html", error="Room code invalid")

        session["room_code"] = room_code
        session["username"] = username
        return redirect(url_for("room"))

    info = request.args.get("info") if request.args.get("info") else ""
    return render_template("home.html", info=info)

def generate_room_code(*, length: int, existing_codes: list[str]) -> str:
    while True:
        code_chars = [random.choice(ascii_letters) for _ in range(length)]
        code = "".join(code_chars)
        if code not in existing_codes:
            return code

@app.route("/new_room")
def new_room():
    import redis
    rd = redis.Redis(host="127.0.0.1", port=6379, db=0)
    room_code = generate_room_code(length=ROOM_CODE_LENGTH, existing_codes=list(rd.keys()))
    new_room = {
        "members": 0,
        "messages": [],
    }
    cache.set(room_code, new_room, timeout=AN_HOUR_IN_SEC)
    return redirect(url_for("home", info=room_code))

@app.route("/room")
def room():
    room_code = session.get("room_code")
    username = session.get("username")

    if username is None or room_code is None or not cache.get(room_code):
        return redirect(url_for("home"))

    room = cache.get(room_code)
    messages = room.get("messages") or []
    return render_template("room.html", room_code=room_code, username=username, messages=messages)


@socketio.on("connect")
def handle_connect():
    username = session.get("username")
    room_code = session.get("room_code")

    if username is None or room_code is None:
        return

    if not cache.get(room_code):
        leave_room(room=room_code)

    join_room(room=room_code)
    message = {
        "sender": "",
        "message": f"{username} joined the chat",
    }
    send(message=message, to=room_code)

    room = cache.get(room_code)
    room.update({"members": room.get("members", 0) + 1})
    cache.set(room_code, room, timeout=AN_HOUR_IN_SEC)

@socketio.on("message")
def handle_message(payload):
    username = session.get("username")
    room_code = session.get("room_code")

    if not cache.get(room_code):
        return

    message = {
        "sender": username,
        "message": payload.get("message"),
    }
    send(message=message, to=room_code)

    room = cache.get(room_code)
    room_messages = room.get("messages") or []
    room_messages.append(message)
    room.update({"messages": room_messages})
    cache.set(room_code, room, timeout=AN_HOUR_IN_SEC)

@socketio.on("disconnect")
def handle_disconnect():
    username = session.get("username")
    room_code = session.get("room_code")

    leave_room(room=room_code)

    if cache.get(room_code):
        room = cache.get(room_code)

        if room.get("members") <= 0:
            cache.delete(room_code)
        else:
            room.update({"members": room.get("members", 0) - 1})
            cache.set(room_code, room, timeout=AN_HOUR_IN_SEC)

        messsge = {
            "sender": "",
            "message": f"{username} has left the chat",
        }
        send(message=messsge, to=room_code)


if __name__ == "__main__":
    socketio.run(app, debug=True)