"""WiFi connection management"""
import network
import time

def connect_wifi(ssid, password):
    """Connect to WiFi network"""
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    
    if not wlan.isconnected():
        print('Connecting to WiFi...')
        wlan.connect(ssid, password)
        while not wlan.isconnected():
            time.sleep(1)
    
    print('Connected to WiFi')
    print('IP:', wlan.ifconfig()[0])
    return wlan