Bug fixes and TODO cleanup