Better error checking and code cleanup