Refacotred code + fixed milestone edit issue