Clean up main.go a bit