from functools import wraps
from flask import redirect, url_for, session


def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if not session.get("logged_in"):
            # If user is not logged in, redirect to login page
            return redirect("/login")
        return f(*args, **kwargs)

    return decorated_function


def team_login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if not session.get("team_logged_in"):
            return redirect("/testPDF/login")
        return f(*args, **kwargs)

    return decorated_function