Refactor: move logic from main to logic class