rework it a bit