Refactored for better error handling