Refactored and fixed edit