refactored some code and fixed a small bug.

Signed-off-by: Chris Jones <chris@jonezy.org>