Fix/add shit