Code refactor to add more types and organize packages