Small refactor and cleanup so I can be a little les ashamed.