Updated code structure and functionality