import rlcd
from machine import Pin, SPI
from wifi_util import WiFiUtil
import time

def main():
    print("Initializing SPI and Display...")
    # Setup SPI
    spi = SPI(1, baudrate=20000000, polarity=0, phase=0, 
              sck=Pin(11), mosi=Pin(12))

    # Initialize Display
    display = rlcd.RLCD(spi, cs=Pin(40), dc=Pin(5), rst=Pin(41))

    # 1. Show a loading state
    display.clear(0)
    display.text("Scanning for Wi-Fi...", 10, 10, 1)
    display.show()

    # 2. Scan for networks
    wifi = WiFiUtil()
    networks = wifi.scan()

    # 3. Draw the results
    print("Drawing networks to buffer...")
    display.clear(0)  # Clear to white
    
    # Title
    display.text("Available Wi-Fi Networks:", 10, 10, 1)
    display.line(10, 20, 210, 20, 1)
    
    # List networks
    y_offset = 30
    # Limit to 20 networks so they fit on the 300px tall screen (20 * 12px = 240px + 30px offset = 270px)
    for i, net in enumerate(networks[:20]):
        text = f"{i+1}. {net['ssid']} ({net['rssi']}dBm)"
        display.text(text, 10, y_offset, 1)
        y_offset += 12

    if not networks:
        display.text("No networks found.", 10, 30, 1)

    # 4. Update Screen
    print("Updating screen...")
    display.show()
    try:
        display.save_screenshot('screenshot.pbm')
    except Exception as e:
        print(f"Failed to save screenshot: {e}")
    print("Done!")

if __name__ == "__main__":
    main()