Fix selection background too

Refactor code to centralize the logic too