[OK] refactor the main fun