Fixed gradle build and removed TODO