Move more logic over
Extract helper and polish UI