Reverting previous commit and fixing it the right way