Clean up the PR.

Use gofmt instead of goimports to not change the order of the imports.
Consolidate TODO comments into a short version.