Improving error handling. Needs to be tested.