try to improve code design