Revert last two commits, fix it right