Cleanup & error handling
Updated version