Refactor

@kenoir please can you review and let me know if you're OK with it and if so merge and I'll write some tests around the code