Add flush and refactor because function was too complex