from machine import UART, Pin, I2C
import time
from umqtt.simple import MQTTClient
from ssd1306 import SSD1306_I2C  # OLED 驱动库

# -------------------- 配置参数 --------------------
WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""
MQTT_CLIENT_ID = "8854b385911048edaa12808fb1f199e6"  # 巴法云 UID
MQTT_BROKER = "bemfa.com"
MQTT_PORT = 9501
MQTT_TOPIC = "myled"        # LED 控制主题（巴法云创建）
MQTT_OLED_TOPIC = "mysta"   # OLED 状态显示主题（巴法云创建）
MQTT_USER = "+8618840167598"  # 巴法云 UID
MQTT_PASSWORD = "8854b385911048edaa12808fb1f199e6"  # 巴法云私钥

# -------------------- 硬件配置 --------------------
LED_PIN = 15  # LED 引脚
UART_ID = 1   # UART 编号
UART_BAUDRATE = 115200
UART_TX = 1   # 发送引脚
UART_RX = 3   # 接收引脚

# OLED 配置（I2C 总线，根据仿真平台调整引脚）
i2c = I2C(0, scl=Pin(22), sda=Pin(21))  # SCL=GPIO22, SDA=GPIO21
oled_width = 128
oled_height = 64
oled = SSD1306_I2C(oled_width, oled_height, i2c)

# 初始化硬件
led = Pin(LED_PIN, Pin.OUT)
uart = UART(UART_ID, baudrate=UART_BAUDRATE, tx=UART_TX, rx=UART_RX)

# -------------------- 核心函数 --------------------
def connect_wifi():
    import network
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        print('Connecting to network...')
        wlan.connect(WIFI_SSID, WIFI_PASSWORD)
        timeout = 10
        while timeout > 0:
            if wlan.isconnected():
                break
            timeout -= 1
            time.sleep(1)
        if not wlan.isconnected():
            raise RuntimeError("WiFi connection failed")
    print('Network config:', wlan.ifconfig())

def connect_mqtt():
    client = MQTTClient(MQTT_CLIENT_ID, MQTT_BROKER, MQTT_PORT, MQTT_USER, MQTT_PASSWORD)
    client.connect()
    print(f"Connected to MQTT broker: {MQTT_BROKER}")
    return client

def publish_led_state(client, state):
    payload = "ON" if state else "OFF"
    try:
        client.publish(MQTT_TOPIC, payload)
        print(f"Published to {MQTT_TOPIC}: {payload}")
        # 同步更新 OLED 显示 + 发布到 OLED 主题
        update_oled(payload)
        client.publish(MQTT_OLED_TOPIC, payload)
        print(f"Published to {MQTT_OLED_TOPIC}: {payload}")
    except Exception as e:
        print(f"MQTT publish error: {e}")

def update_oled(state_str):
    """更新 OLED 显示内容"""
    oled.fill(0)  # 清屏
    oled.text("LED State:", 0, 10)
    oled.text(state_str, 0, 30)
    oled.show()   # 刷新显示

# -------------------- 主程序 --------------------
def main():
    try:
        # 连接 WiFi 和 MQTT
        connect_wifi()
        mqtt_client = connect_mqtt()

        # 初始化提示
        uart.write("\r\n=== LED Control ===\r\n")
        uart.write("Enter '1' to turn LED ON\r\n")
        uart.write("Enter '0' to turn LED OFF\r\n")
        uart.write("-------------------\r\n")

        # 主循环
        while True:
            uart.write("> ")  # 显示提示符
            time.sleep(0.1)

            # 等待用户输入（5 秒超时）
            start_time = time.ticks_ms()
            command = b""
            while (time.ticks_ms() - start_time) < 5000:  # 5 秒超时
                if uart.any():
                    char = uart.read(1)  # 读取一个字符
                    if char == b'\r':  # 回车结束输入
                        break
                    command += char
                    uart.write(char)  # 回显输入
                    time.sleep(0.01)

            # 处理命令
            command = command.strip()
            if command == b'1':
                led.on()
                uart.write("\r\nLED is now ON\r\n")
                publish_led_state(mqtt_client, True)
            elif command == b'0':
                led.off()
                uart.write("\r\nLED is now OFF\r\n")
                publish_led_state(mqtt_client, False)
            elif command:  # 非空但无效命令
                uart.write(f"\r\nInvalid command: {command.decode()}\r\n")
            else:  # 超时无输入
                uart.write("\r\nInput timed out\r\n")

    except KeyboardInterrupt:
        print("Program terminated by user")
    except Exception as e:
        print(f"Fatal error: {e}")
    finally:
        # 清理资源
        if 'mqtt_client' in locals():
            mqtt_client.disconnect()
        led.off()
        print("Resources cleaned up")

# 执行主程序
if __name__ == "__main__":
    main()