Refactor some things in main.go