Heavy refactoring for better separation of logic