from machine import UART, Pin, SoftI2C
import network
import time
from umqtt.simple import MQTTClient
from ssd1306 import SSD1306_I2C

# ------------------ 配置参数 ------------------ #
# WiFi 配置
WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""

# 巴法云 MQTT 配置
BAFA_CLIENT_ID = "bd6ff39ba76075f7d2dfbc26e14658c7"
BAFA_TOPIC = "light002"
MQTT_BROKER = "bemfa.com"
MQTT_PORT = 9501

# OLED 屏幕配置
i2c = SoftI2C(sda=Pin(23), scl=Pin(18))
oled_width = 128
oled_height = 64
oled = SSD1306_I2C(oled_width, oled_height, i2c)

# ------------------ 硬件初始化 ------------------ #
uart = UART(1, baudrate=115200, tx=1, rx=3)
led = Pin(15, Pin.OUT)
led_status = "off"  # 记录 LED 状态

# MQTT 客户端对象（全局）
mqtt_client = None

# ------------------ 函数定义 ------------------ #
# 连接 WiFi
def connect_wifi():
    uart.write(b"\rConnecting to WiFi...\n")
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    
    if not wlan.isconnected():
        wlan.connect(WIFI_SSID, WIFI_PASSWORD)
        max_wait = 20
        while max_wait > 0:
            if wlan.isconnected():
                break
            max_wait -= 1
            time.sleep(1)
            uart.write(b".")
    
    if wlan.isconnected():
        uart.write(b"\rWiFi connected! IP: " + wlan.ifconfig()[0] + b"\n")
        return True
    else:
        uart.write(b"\rWiFi failed!\n")
        return False

# 连接 MQTT 服务器
def connect_mqtt():
    global mqtt_client
    uart.write(b"\rConnecting to MQTT Broker...\n")
    
    mqtt_client = MQTTClient(
        client_id=BAFA_CLIENT_ID, 
        server=MQTT_BROKER, 
        port=MQTT_PORT,
        keepalive=60
    )
    
    try:
        mqtt_client.connect()
        uart.write(b"\rMQTT connected! Ready to send/receive.\n")
        return True
    except Exception as e:
        uart.write(b"\rMQTT connect failed: " + str(e).encode() + b"\n")
        return False

# MQTT 消息回调
def on_mqtt_message(topic, msg):
    uart.write(b"\rReceived MQTT message: " + topic + b" -> " + msg + b"\n")
    # 只控制 LED，不更新 OLED（避免自动反转显示）
    if msg == b"on":
        led.value(1)
        global led_status
        led_status = "on"
        uart.write(b"\rLED ON (MQTT command)\n")
    elif msg == b"off":
        led.value(0)
        led_status = "off"
        uart.write(b"\rLED OFF (MQTT command)\n")
    # 不调用 control_led()，避免触发 OLED 更新

# 控制 LED 并同步到巴法云（仅用于串口手动输入）
def control_led(status):
    global led_status
    if status == "on":
        led.value(1)
        led_status = "on"
        uart.write(b"\rLED ON (Input 1)\n")
    elif status == "off":
        led.value(0)
        led_status = "off"
        uart.write(b"\rLED OFF (Input 0)\n")
    else:
        uart.write(b"\rInvalid input! Use '1' or '0'.\n")
        return
    
    # 发布状态到巴法云
    mqtt_client.publish(BAFA_TOPIC, led_status)
    uart.write(b"\rSent to Bemfa: " + led_status.encode() + b"\n")
    
    # 更新 OLED 显示（仅在手动输入时更新）
    oled.fill(0)
    oled.text("LED Status: " + led_status, 0, 0)
    oled.show()

# 主程序逻辑
def main():
    # 1. 连接 WiFi
    if not connect_wifi():
        return
    
    # 2. 连接 MQTT 服务器
    if not connect_mqtt():
        return
    
    # 3. 订阅主题
    mqtt_client.set_callback(on_mqtt_message)
    mqtt_client.subscribe(BAFA_TOPIC)
    uart.write(b"\rSubscribed to topic: " + BAFA_TOPIC + b"\n")
    
    # 4. 初始化 OLED 显示
    oled.fill(0)
    oled.text("Ready for input...", 0, 0)
    oled.show()
    
    # 5. 主循环
    uart.write(b"\rSystem ready! Enter '1'/'0' to control LED.\n")
    while True:
        # 检查串口输入（仅在这里处理 OLED 更新）
        if uart.any():
            data = uart.read().strip()
            if data == b"1":
                control_led("on")
            elif data == b"0":
                control_led("off")
            else:
                uart.write(b"\rInvalid input! Try '1' (ON) or '0' (OFF).\n")
        
        # 检查 MQTT 消息
        if mqtt_client is not None:
            mqtt_client.check_msg()
        
        time.sleep(0.1)

# ------------------ 启动程序 ------------------ #
if __name__ == "__main__":
    main()