Fixed null check and organized imports