refactor main py