Main.java quick update! fixed some possible code duplication