Code cleanup - applied better class abstraction & added more helper methods