Some code review, small fixes and TODO comments