Refactored a bit and used interfaces