Updated main function