Rework some stuff :)