from machine import Pin, I2C, freq
import ssd1306
import time
import network
import gc
import os
import sdcard

# Wi-Fi credentials
SSID = 'your_SSID'  # Replace with your Wi-Fi SSID
PASSWORD = 'your_PASSWORD'  # Replace with your Wi-Fi password

# Initialize I2C for OLED
i2c = I2C(0, scl=Pin(22), sda=Pin(21), freq=50000)  # Set frequency to 50 kHz

# Scan for I2C devices
devices = i2c.scan()
if devices:
    print("I2C devices found:", devices)
else:
    print("No I2C devices found")

# Initialize OLED if found
if 0x3c in devices:  # Check for the typical SSD1306 address
    oled_width = 128
    oled_height = 64
    oled = ssd1306.SSD1306_I2C(oled_width, oled_height, i2c)

    # Clear the display and show initial message
    oled.fill(0)  # Clear the display
    oled.text('ESP32 Stats:', 0, 0)  # Display header
    oled.show()  # Refresh the display

    # Function to connect to Wi-Fi
    def connect_wifi():
        wlan = network.WLAN(network.STA_IF)
        wlan.active(True)
        wlan.connect(SSID, PASSWORD)
        
        # Wait for connection
        timeout = 10  # seconds
        while not wlan.isconnected() and timeout > 0:
            timeout -= 1
            time.sleep(1)
        
        if wlan.isconnected():
            print('Connected to Wi-Fi:', wlan.ifconfig())
        else:
            print('Failed to connect to Wi-Fi.')

    # Function to initialize the SD card and return free space
    def init_sd():
        sd = sdcard.SDCard(spi=Pin(23), cs=Pin(5))  # Adjust SPI and CS pins as needed
        os.mount(sd, '/sd')
        free_space = os.statvfs('/sd')[0] * os.statvfs('/sd')[3]  # Total space in bytes
        return free_space

    # Connect to Wi-Fi
    connect_wifi()

    # Initialize SD card and get free space
    try:
        sd_space = init_sd()
        print('SD Card initialized. Free space:', sd_space, 'bytes')
    except Exception as e:
        print('Failed to initialize SD card:', e)
        sd_space = 0  # Set to 0 if the SD card cannot be initialized

    while True:
        # Clear the display for new data
        oled.fill(0)

        # Display ESP32 statistics
        oled.text('CPU Freq: {} MHz'.format(freq() // 1000000), 0, 10)
        oled.text('Free Mem: {} bytes'.format(gc.mem_free()), 0, 20)

        # Check Wi-Fi status
        wlan = network.WLAN(network.STA_IF)
        if wlan.isconnected():
            ip = wlan.ifconfig()[0]
            oled.text('IP Address: {}'.format(ip), 0, 30)
        else:
            oled.text('WiFi: Not connected', 0, 30)

        # Display SD card space
        oled.text('SD Free: {} bytes'.format(sd_space), 0, 40)

        # Show the updated display
        oled.show()

        # Delay for readability
        time.sleep(2)
else:
    print("OLED display not found!")