Updated Code with more functionality