Total code refactor thanks to IntelliJ