from cryptography.fernet import Fernet
from telegram import Update
from dbstore.peewee_store import db, User, get_user_by_id
import functools


def encrypt(input: str, key: str) -> str:
    f = Fernet(key)
    return f.encrypt(bytes(input, 'utf-8')).decode('utf-8')


def decrypt(input: str, key: str) -> str:
    f = Fernet(key)
    return f.decrypt(input).decode('utf-8')


def check_user(fn):
    """
    Decorator: loads User model and passes it to the function or stops the request.
    Ref: https://shallowdepth.online/posts/2021/12/using-python-decorators-to-process-and-authorize-requests/
    """

    @functools.wraps(fn)
    async def wrapper(*args, **kwargs):
        # Expects that Update object is always the first arg
        update: Update = args[0]
        user = get_user_by_id(str(update.effective_user.id))

        if len(user) == 0:
            await update.effective_message.reply_text("You are not logged in. Use `/login` to link your account first")
            return
        else:
            # TODO
            # check access_key is still valid
            pass
        return await fn(*args, **kwargs, user=user)

    return wrapper


# if __name__ == "__main__":
#     key = Fernet.generate_key().decode('utf-8')
#     print(key)
#     print(encrypt("Hello World!", key))
#     print(decrypt(encrypt("Hello World!", key), key))