Add update functionality
Refactor and add comments