Little refactor of main.lua