Updated with more checks