Refactor for better organization