Improve code and increase test coverage

We have identified multiple areas in the code which needs to be refactored to make them more maintainable and scalable in the future. We have also observed that we need to increase unit testing so that we can make sure there is no regression when making modifications to this module.