Refactor code according new requirements