Code cleanup and refactor for consistency