[TASK] Refactor to simplify code