Cleaning some things up, adding some error handling