Half done refactor of code...