Just some cleanup, this requires refactoring though