import os
import gc
import wifi
import socketpool
import ssl
import time
import rtc
import adafruit_requests
from adafruit_io.adafruit_io import IO_HTTP
import displayio
import board
import supervisor
from setup_display import setup_display
from adafruit_display_shapes.rect import Rect
from adafruit_display_shapes.roundrect import RoundRect
from adafruit_display_shapes.line import Line
from adafruit_bitmap_font import bitmap_font
from adafruit_display_text import label, wrap_text_to_pixels

# Debug flag - set to False to reduce memory usage
DEBUG = False

# Global list to store health texts for rotation (text, color)
health_texts = []

def format_pollutant(value):
    """Helper function to format pollutant values"""
    return " - - " if value is None else (f"{value} ug" if value != 'N/A' else 'N/A')

def hex_to_int(color_str):
    """Convert hex color string to integer"""
    if color_str and color_str != 'N/A' and color_str.startswith('#'):
        return int(color_str.replace('#', '0x'), 16)
    return 0xFF0000  # Default to red if invalid

def title_case(s):
    # Replace underscores with spaces first
    s = s.replace('_', ' ')
    if len(s) > 3:
        words = s.split()
        titled = [word[0].upper() + word[1:].lower() if len(word) > 0 else '' for word in words]
        return ' '.join(titled)
    else:
        return s

def update_display_with_data(data_aqio, data_aqie):
    """
    Update display with AQI data from both feeds
    """
    global health_texts
    
    # Update last update time from RTC
    current_time = time.localtime()
    
    # Collect health texts from both sources (alternating between services)
    health_texts = []
    
    # Update AQIO column (center)
    if data_aqio:
        line1_c.text = title_case(data_aqio.get('provider', 'N/A'))
        current = data_aqio.get('current', {})
        
        line2_c.text = str(current.get('aqi', 'N/A'))
        line5_c.text = current.get('category_label', 'N/A')
        line8_c.text = format_pollutant(current.get('pm2_5'))
        line9_c.text = format_pollutant(current.get('o3'))
        line10_c.text = (current.get('primary_pollutant', 'N/A') or 'N/A').upper()
        line11_c.text = format_pollutant(current.get('pm10'))
        line12_c.text = format_pollutant(current.get('co'))
        line13_c.text = format_pollutant(current.get('no2'))
        line14_c.text = format_pollutant(current.get('so2'))
        o_color.fill = hex_to_int(current.get('category_color', '#FF0000'))
        
        forecast_today = data_aqio.get('forecast_today', {})
        line3_c.text = str(forecast_today.get('aqi', 'N/A'))
        line6_c.text = forecast_today.get('category_label', 'N/A')
        o_color_f.fill = hex_to_int(forecast_today.get('category_color', '#FF0000'))
        
        forecast_tomorrow = data_aqio.get('forecast_tomorrow', {})
        line4_c.text = str(forecast_tomorrow.get('aqi', 'N/A'))
        line7_c.text = forecast_tomorrow.get('category_label', 'N/A')
        o_color_t.fill = hex_to_int(forecast_tomorrow.get('category_color', '#FF0000'))
        
        # Collect OM health texts with time period prefixes and colors
        health_texts.append((f"Now: {current.get('health_description', 'N/A')}", 0xFFFFFF, 0))  # White
        health_texts.append((f"Now: {current.get('health_recommendation', 'N/A')}", 0xFFFFFF, 1))  # White
        health_texts.append((f"Today: {forecast_today.get('health_description', 'N/A')}", 0xFFFF00, 4))  # Yellow
        health_texts.append((f"Today: {forecast_today.get('health_recommendation', 'N/A')}", 0xFFFF00, 5))  # Yellow
        health_texts.append((f"Tomorrow: {forecast_tomorrow.get('health_description', 'N/A')}", 0x00FFFF, 8))  # Cyan
        health_texts.append((f"Tomorrow: {forecast_tomorrow.get('health_recommendation', 'N/A')}", 0x00FFFF, 9))  # Cyan
    
    # Update AQIE column (right)
    if data_aqie:
        line1_r.text = title_case(data_aqie.get('provider', 'N/A'))
        current = data_aqie.get('current', {})
        
        line2_r.text = str(current.get('aqi', 'N/A'))
        line5_r.text = current.get('category_label', 'N/A')
        line8_r.text = format_pollutant(current.get('pm2_5'))
        line9_r.text = format_pollutant(current.get('o3'))
        line10_r.text = (current.get('primary_pollutant', 'N/A') or 'N/A').upper()
        line11_r.text = format_pollutant(current.get('pm10'))
        line12_r.text = format_pollutant(current.get('co'))
        line13_r.text = format_pollutant(current.get('no2'))
        line14_r.text = format_pollutant(current.get('so2'))
        e_color.fill = hex_to_int(current.get('category_color', '#FF0000'))
        
        forecast_today = data_aqie.get('forecast_today', {})
        line3_r.text = str(forecast_today.get('aqi', 'N/A'))
        line6_r.text = forecast_today.get('category_label', 'N/A')
        e_color_f.fill = hex_to_int(forecast_today.get('category_color', '#FF0000'))
        
        forecast_tomorrow = data_aqie.get('forecast_tomorrow', {})
        line4_r.text = str(forecast_tomorrow.get('aqi', 'N/A'))
        line7_r.text = forecast_tomorrow.get('category_label', 'N/A')
        e_color_t.fill = hex_to_int(forecast_tomorrow.get('category_color', '#FF0000'))
        
        # Collect Airnow health texts with time period prefixes and colors
        health_texts.append((f"Now: {current.get('health_description', 'N/A')}", 0xFFFFFF, 2))  # White
        health_texts.append((f"Now: {current.get('health_recommendation', 'N/A')}", 0xFFFFFF, 3))  # White
        health_texts.append((f"Today: {forecast_today.get('health_description', 'N/A')}", 0xFFFF00, 6))  # Yellow
        health_texts.append((f"Today: {forecast_today.get('health_recommendation', 'N/A')}", 0xFFFF00, 7))  # Yellow
        health_texts.append((f"Tomorrow: {forecast_tomorrow.get('health_description', 'N/A')}", 0x00FFFF, 10))  # Cyan
        health_texts.append((f"Tomorrow: {forecast_tomorrow.get('health_recommendation', 'N/A')}", 0x00FFFF, 11))  # Cyan
    
    # Sort health_texts by the order integer (third element)
    health_texts.sort(key=lambda x: x[2])

def get_aqi_data(feed_key, timeout=30):
    """
    Retrieve AQI data from Adafruit IO feed
    Args:
        feed_key: Feed key string ('aquo' or 'aque')
        timeout: Request timeout in seconds (default: 30)
    Returns:
        JSON data from the feed
    """
    try:
        print(f"\nRetrieving AQI data from feed: {feed_key}")
        # Force garbage collection before network request
        gc.collect()
        data = io.receive_aqi(feed_key)
        parse_aqi_data(data)
        return data
    except OSError as e:
        # Network-related errors (ETIMEDOUT, ECONNABORTED, etc.)
        error_code = getattr(e, 'errno', None)
        if error_code == 116:
            print(f"Connection timeout (ETIMEDOUT) for feed {feed_key}")
        elif error_code == 113:
            print(f"Connection aborted (ECONNABORTED) for feed {feed_key}")
        else:
            print(f"Network error retrieving AQI data: {e}")
        return None
    except Exception as e:
        print(f"Error retrieving AQI data: {e}")
        return None

def parse_aqi_data(data):
    """
    Parse and display AQI data from JSON response
    Args:
        data: JSON data dictionary from Adafruit IO
    """
    if not data or not DEBUG:
        return

    print("\n" + "=" * 60)
    print(f"AIR QUALITY DATA - {data.get('name', 'Unknown')}")
    print(f"Provider: {data.get('provider', 'Unknown').upper()}")
    print("=" * 60)

    # Parse current conditions
    if 'current' in data:
        current = data['current']
        print("\nCURRENT CONDITIONS:")
        print(f"  AQI: {current.get('aqi', 'N/A')}")
        print(f"  Category: {current.get('category_label', 'N/A')}")
        print(f"  Color: {current.get('category_color', 'N/A')}")
        print(f"  Timestamp: {current.get('timestamp', 'N/A')}")

        # Show location info if available (EPA/AirNow data)
        if current.get('reporting_area'):
            print(f"  Location: {current.get('reporting_area', '')}, {current.get('state', '')}")

        print(f"\n  POLLUTANTS:")
        # Only show pollutants that have values
        if current.get('pm2_5') is not None:
            print(f"    PM2.5: {current.get('pm2_5')} µg/m³")
        if current.get('pm10') is not None:
            print(f"    PM10:  {current.get('pm10')} µg/m³")
        if current.get('co') is not None:
            print(f"    CO:    {current.get('co')} µg/m³")
        if current.get('o3') is not None:
            print(f"    O3:    {current.get('o3')} µg/m³")
        if current.get('no2') is not None:
            print(f"    NO2:   {current.get('no2')} µg/m³")
        if current.get('so2') is not None:
            print(f"    SO2:   {current.get('so2')} µg/m³")
        if current.get('primary_pollutant'):
            print(f"    Primary: {current.get('primary_pollutant').upper()}")

        print(f"\n  HEALTH:")
        print(f"    {current.get('health_recommendation', 'N/A')}")

    # Parse forecast today
    if 'forecast_today' in data:
        forecast = data['forecast_today']
        print(f"\nTODAY'S FORECAST:")
        print(f"  AQI: {forecast.get('aqi', 'N/A')}")
        print(f"  Category: {forecast.get('category_label', 'N/A')}")
        if forecast.get('primary_pollutant'):
            print(f"  Primary Pollutant: {forecast.get('primary_pollutant').upper()}")
        print(f"  Recommendation: {forecast.get('health_recommendation', 'N/A')}")

    # Parse forecast tomorrow
    if 'forecast_tomorrow' in data:
        forecast = data['forecast_tomorrow']
        print(f"\nTOMORROW'S FORECAST:")
        print(f"  AQI: {forecast.get('aqi', 'N/A')}")
        print(f"  Category: {forecast.get('category_label', 'N/A')}")
        if forecast.get('primary_pollutant'):
            print(f"  Primary Pollutant: {forecast.get('primary_pollutant').upper()}")
        print(f"  Recommendation: {forecast.get('health_recommendation', 'N/A')}")

    print("\n" + "=" * 60)

setup_display()
display = supervisor.runtime.display

# Load fonts
font_11 = bitmap_font.load_font("fonts/ArialMT-11.pcf")
font_12 = bitmap_font.load_font("fonts/ArialMT-12.pcf")
print("Fonts loaded successfully")

# Create display group and add navy blue background
main_group = displayio.Group()
navy_blue = 0x000066  # Navy blue color
background = Rect(0, 0, display.width, display.height, fill=navy_blue)
main_group.append(background)

# Create text labels - Left column (data point names)
text_color = 0xFFFFFF  # White text

cu_time = label.Label(font_11, text="Current time:", color=0x00FFFF, anchored_position = (2, 2), anchor_point=(0.0, 0.0))
main_group.append(cu_time)

line2 = label.Label(font_12, text="Current AQI:", color=text_color, anchored_position = (100, 25), anchor_point=(1.0, 0.0))
main_group.append(line2)

line3 = label.Label(font_12, text="Forecast AQI:", color=0xFFFF00, anchored_position = (100, 40), anchor_point=(1.0, 0.5))
main_group.append(line3)

line4 = label.Label(font_12, text="Tomorrow AQI:", color=0x00FFFF, anchored_position = (100, 55), anchor_point=(1.0, 1.0))
main_group.append(line4)

line5 = label.Label(font_12, text="Current Cat:", color=text_color, anchored_position = (100, 68), anchor_point=(1.0, 0.0))
main_group.append(line5)

line6 = label.Label(font_12, text="Forecast Cat:", color=0xFFFF00, anchored_position = (100, 83), anchor_point=(1.0, 0.5))
main_group.append(line6)

line7 = label.Label(font_12, text="Tomorrow Cat:", color=0x00FFFF, anchored_position = (100, 98), anchor_point=(1.0, 1.0))
main_group.append(line7)

line8 = label.Label(font_11, text="PM2.5:", color=text_color, anchored_position = (100, 113), anchor_point=(1.0, 0.0))
main_group.append(line8)

line9 = label.Label(font_11, text="O3:", color=text_color, anchored_position = (100, 125), anchor_point=(1.0, 0.0))
main_group.append(line9)

line10 = label.Label(font_11, text="Primary Pollutant:", color=0xFFFF00, anchored_position = (100, 137), anchor_point=(1.0, 0.0))
main_group.append(line10)

line11 = label.Label(font_11, text="PM10:", color=text_color, anchored_position = (100, 149), anchor_point=(1.0, 0.0))
main_group.append(line11)

line12 = label.Label(font_11, text="CO:", color=text_color, anchored_position = (100, 161), anchor_point=(1.0, 0.0))
main_group.append(line12)

line13 = label.Label(font_11, text="NO2:", color=text_color, anchored_position = (100, 173), anchor_point=(1.0, 0.0))
main_group.append(line13)

line14 = label.Label(font_11, text="SO2:", color=text_color, anchored_position = (100, 185), anchor_point=(1.0, 0.0))
main_group.append(line14)

line15 = label.Label(font_11, text="text:", color=text_color, anchored_position = (10, 199), anchor_point=(0.0, 0.0))
main_group.append(line15)

# Create center column
line1_c = label.Label(font_12, text="Line 1_c", color=0x44FF44, anchored_position = (120, 10), anchor_point=(0.0, 0.5))
main_group.append(line1_c)

line2_c = label.Label(font_12, text="Line 2_c", color=text_color, anchored_position = (120, 25), anchor_point=(0.0, 0.0))
main_group.append(line2_c)

line3_c = label.Label(font_12, text="Line 3_c", color=0xFFFF00, anchored_position = (120, 40), anchor_point=(0.0, 0.5))
main_group.append(line3_c)

line4_c = label.Label(font_12, text="Line 4_c", color=0x00FFFF, anchored_position = (120, 55), anchor_point=(0.0, 1.0))
main_group.append(line4_c)


line5_c = label.Label(font_12, text="Line 5_c", color=text_color, anchored_position = (120, 68), anchor_point=(0.0, 0.0))
main_group.append(line5_c)

line6_c = label.Label(font_12, text="Line 6_c", color=0xFFFF00, anchored_position = (120, 83), anchor_point=(0.0, 0.5))
main_group.append(line6_c)

line7_c = label.Label(font_12, text="Line 7_c", color=0x00FFFF, anchored_position = (120, 98), anchor_point=(0.0, 1.0))
main_group.append(line7_c)

line8_c = label.Label(font_11, text="Line 8_c", color=text_color, anchored_position = (120, 113), anchor_point=(0.0, 0.0))
main_group.append(line8_c)

line9_c = label.Label(font_11, text="Line 9_c", color=text_color, anchored_position = (120, 125), anchor_point=(0.0, 0.0))
main_group.append(line9_c)

line10_c = label.Label(font_11, text="Line 10_c", color=0xFFFF00, anchored_position = (120, 137), anchor_point=(0.0, 0.0))
main_group.append(line10_c)

line11_c = label.Label(font_11, text="Line 11_c", color=text_color, anchored_position = (120, 149), anchor_point=(0.0, 0.0))
main_group.append(line11_c)

line12_c = label.Label(font_11, text="Line 12_c", color=text_color, anchored_position = (120, 161), anchor_point=(0.0, 0.0))
main_group.append(line12_c)

line13_c = label.Label(font_11, text="Line 13_c", color=text_color, anchored_position = (120, 173), anchor_point=(0.0, 0.0))
main_group.append(line13_c)

line14_c = label.Label(font_11, text="Line 14_c", color=text_color, anchored_position = (120, 185), anchor_point=(0.0, 0.0))
main_group.append(line14_c)


line1_r = label.Label(font_12, text="Line 1_r", color=0x44FF44, anchored_position = (220, 10), anchor_point=(0.0, 0.5))
main_group.append(line1_r)

line2_r = label.Label(font_12, text="Line 2_r", color=text_color, anchored_position = (220, 25), anchor_point=(0.0, 0.0))
main_group.append(line2_r)

line3_r = label.Label(font_12, text="Line 3_r", color=0xFFFF00, anchored_position = (220, 40), anchor_point=(0.0, 0.5))
main_group.append(line3_r)

line4_r = label.Label(font_12, text="Line 4_r", color=0x00FFFF, anchored_position = (220, 55), anchor_point=(0.0, 1.0))
main_group.append(line4_r)

line5_r = label.Label(font_12, text="Line 5_r", color=text_color, anchored_position = (220, 68), anchor_point=(0.0, 0.0))
main_group.append(line5_r)

line6_r = label.Label(font_12, text="Line 6_r", color=0xFFFF00, anchored_position = (220, 83), anchor_point=(0.0, 0.5))
main_group.append(line6_r)

line7_r = label.Label(font_12, text="Line 7_r", color=0x00FFFF, anchored_position = (220, 98), anchor_point=(0.0, 1.0))
main_group.append(line7_r)

line8_r = label.Label(font_11, text="Line 8_r", color=text_color, anchored_position = (220, 113), anchor_point=(0.0, 0.0))
main_group.append(line8_r)

line9_r = label.Label(font_11, text="Line 9_r", color=text_color, anchored_position = (220, 125), anchor_point=(0.0, 0.0))
main_group.append(line9_r)

line10_r = label.Label(font_11, text="Line 10_r", color=0xFFFF00, anchored_position = (220, 137), anchor_point=(0.0, 0.0))
main_group.append(line10_r)

line11_r = label.Label(font_11, text="Line 11_r", color=text_color, anchored_position = (220, 149), anchor_point=(0.0, 0.0))
main_group.append(line11_r)

line12_r = label.Label(font_11, text="Line 12_r", color=text_color, anchored_position = (220, 161), anchor_point=(0.0, 0.0))
main_group.append(line12_r)

line13_r = label.Label(font_11, text="Line 13_r", color=text_color, anchored_position = (220, 173), anchor_point=(0.0, 0.0))
main_group.append(line13_r)

line14_r = label.Label(font_11, text="Line 14_r", color=text_color, anchored_position = (220, 185), anchor_point=(0.0, 0.0))
main_group.append(line14_r)

# Add horizontal separator lines
separator_lineb = Line(0, 62, display.width, 62, color=0x0000FF)
main_group.append(separator_lineb)

separator_linea = Line(0, 20, display.width, 20, color=0x0000FF)
main_group.append(separator_linea)

separator_linec = Line(0, 107, display.width, 107, color=0x0000FF)
main_group.append(separator_linec)

separator_lined = Line(0, 196, display.width, 196, color=0x0000FF)
main_group.append(separator_lined)

# Add colored rectangles for category indicators
o_color = RoundRect(160, 27, 29, 9, r=3, fill=0xFF0000)
main_group.append(o_color)

e_color = RoundRect(260, 27, 29, 9, r=3, fill=0xFF0000)
main_group.append(e_color)

o_color_f = RoundRect(160, 37, 29, 9, r=3, fill=0xFF0000)
main_group.append(o_color_f)

e_color_f = RoundRect(260, 37, 29, 9, r=3, fill=0xFF0000)
main_group.append(e_color_f)

o_color_t = RoundRect(160, 47, 29, 9, r=3, fill=0xFF0000)
main_group.append(o_color_t)

e_color_t = RoundRect(260, 47, 29, 9, r=3, fill=0xFF0000)
main_group.append(e_color_t)

display.root_group = main_group

# Get WiFi credentials from settings.toml
ssid = os.getenv("CIRCUITPY_WIFI_SSID")
password = os.getenv("CIRCUITPY_WIFI_PASSWORD")

aqio = os.getenv("AQIO")
aqie = os.getenv("AQIE")

# Cache settings
aio_username = os.getenv("AIO_USERNAME")
aio_key = os.getenv("AIO_KEY")
timezone = os.getenv("TIMEZONE")

print(f"Connecting to WiFi: {ssid}")

# Connect to WiFi with retry logic
wifi_connected = False
max_retries = 3
for attempt in range(max_retries):
    try:
        wifi.radio.connect(ssid, password)
        print(f"Connected to {ssid}")
        print(f"IP Address: {wifi.radio.ipv4_address}")
        wifi_connected = True
        break
    except Exception as e:
        print(f"WiFi connection attempt {attempt + 1}/{max_retries} failed: {e}")
        if attempt < max_retries - 1:
            time.sleep(5)
        else:
            print("Failed to connect to WiFi after multiple attempts")
            time.sleep(60)
            supervisor.reload()

if wifi_connected:
    # Create socket pool for network requests
    pool = socketpool.SocketPool(wifi.radio)
    print("WiFi connection successful!")

    # Initialize Adafruit IO with SSL context and timeout
    ssl_context = ssl.create_default_context()
    requests = adafruit_requests.Session(pool, ssl_context)
    
    io = IO_HTTP(aio_username, aio_key, requests)
    print(f"Connected to Adafruit IO as: {aio_username}")

    # Initialize RTC and set time from Adafruit IO
    print(f"\nSetting RTC time using timezone: {timezone}")
    try:
        aio_time = io.receive_time()
        if DEBUG:
            print(f"Received time from Adafruit IO: {aio_time}")
        
        rtc_instance = rtc.RTC()
        rtc_instance.datetime = aio_time
        
        current_time = time.localtime()
        print(f"RTC set to: {current_time.tm_year}-{current_time.tm_mon:02d}-{current_time.tm_mday:02d} "
              f"{current_time.tm_hour:02d}:{current_time.tm_min:02d}:{current_time.tm_sec:02d}")
    except Exception as e:
        print(f"Error setting RTC: {e}")

# Retrieve data from both feeds with retry logic
data_aqio = None
data_aqie = None

def check_wifi_connection():
    """Check if WiFi is still connected, attempt reconnect if not"""
    if not wifi.radio.connected:
        print("WiFi disconnected! Attempting to reconnect...")
        try:
            wifi.radio.connect(ssid, password)
            print(f"Reconnected to {ssid}")
            return True
        except Exception as e:
            print(f"WiFi reconnection failed: {e}")
            return False
    return True

def fetch_aqi_with_retry(feed_key, max_attempts=3):
    """Fetch AQI data with exponential backoff retry logic"""
    base_delay = 5  # Start with 5 second delay
    
    for attempt in range(max_attempts):
        # Check WiFi connection before attempting request
        if not check_wifi_connection():
            print("WiFi not connected, skipping attempt")
            if attempt < max_attempts - 1:
                time.sleep(base_delay * (2 ** attempt))
            continue
        
        data = get_aqi_data(feed_key)
        if data:
            print(f"Successfully retrieved data from feed {feed_key}")
            return data
        
        if attempt < max_attempts - 1:
            # Exponential backoff: 5s, 10s, 20s...
            delay = base_delay * (2 ** attempt)
            print(f"Retrying feed {feed_key} in {delay} seconds (attempt {attempt + 1}/{max_attempts})...")
            time.sleep(delay)
        else:
            print(f"Failed to retrieve data from feed {feed_key} after {max_attempts} attempts")
    
    return None

if wifi_connected:
    if aqio:
        data_aqio = fetch_aqi_with_retry(aqio)
    if aqie:
        data_aqie = fetch_aqi_with_retry(aqie)
    
    # Update the display with the retrieved data
    update_display_with_data(data_aqio, data_aqie)
    gc.collect()
    if DEBUG:
        print("gc.mem_free:", gc.mem_free())
print(gc.mem_free())
print("-" * 20)

# Main loop - non-blocking update every hour
if wifi_connected:
    UPDATE_INTERVAL = 3600  # Update every hour (3600 seconds)
    TIME_UPDATE_INTERVAL = 30  # Update current time display every 30 seconds
    HEALTH_TEXT_INTERVAL = 10  # Rotate health text every 10 seconds
    last_aqi_update = time.monotonic()
    last_clock_update = time.monotonic()
    last_health_update = time.monotonic()
    health_text_index = 0
    
    print("\nEntering main loop...")
    
    while True:
        now = time.monotonic()
        
        # Update current time display every 30 seconds
        if now - last_clock_update >= TIME_UPDATE_INTERVAL:
            rtc_time = time.localtime()
            # Calculate minutes until next data update
            elapsed = now - last_aqi_update
            remaining_seconds = UPDATE_INTERVAL - elapsed
            remaining_minutes = int(remaining_seconds / 60)
            cu_time.text = f"{rtc_time.tm_hour:02d}:{rtc_time.tm_min:02d} ({remaining_minutes})"
            last_clock_update = now
        
        # Rotate health text every 10 seconds
        if now - last_health_update >= HEALTH_TEXT_INTERVAL and health_texts:
            # Get current health text, color, and order
            text, color, order = health_texts[health_text_index]
            wrapped_text = "\n".join(wrap_text_to_pixels(text, 310, font_11))
            line15.text = wrapped_text
            line15.color = color
            
            # Move to next text
            health_text_index = (health_text_index + 1) % len(health_texts)
            last_health_update = now
        
        # Check if it's time to update AQI data (1 hour has passed)
        if now - last_aqi_update >= UPDATE_INTERVAL:
            print("\nUpdating AQI data...")
            
            # Verify WiFi is still connected before attempting update
            if not check_wifi_connection():
                print("WiFi connection lost, skipping this update cycle")
                last_aqi_update = now  # Reset timer to try again in next cycle
            else:
                # Retrieve updated data with retry
                if aqio:
                    data_aqio = fetch_aqi_with_retry(aqio)
                if aqie:
                    data_aqie = fetch_aqi_with_retry(aqie)
                
                # Update the display
                update_display_with_data(data_aqio, data_aqie)
                health_text_index = 0  # Reset to first text after update
                gc.collect()
                if DEBUG:
                    print("gc.mem_free:", gc.mem_free())
                
                # Reset the timer
                last_aqi_update = now
        
        # Short sleep to prevent busy-waiting
        time.sleep(1)
else:
    print("WiFi not connected. Reloading system...")
    time.sleep(60)
    supervisor.reload()