Refactor and modularize