Adding exception handling and cleaning the code