Move logic to model