Ading some TODOs in