import machine
import os
import time
import network
from umqtt.simple import MQTTClient
import ssd1306

WIFI_SSID = "Wokwi-GUEST"
WIFI_PASS = ""
CHANNEL_ID = "3338011"
MQTT_CLIENT_ID = "GQoXNioWECA1DwYSFxgHCig"
MQTT_USER = "GQoXNioWECA1DwYSFxgHCig"
MQTT_PASS = "EkR0v1OBvZRmmpeVHY4QgNol"
MQTT_HOST = "mqtt3.thingspeak.com"
TOPIC = "channels/" + CHANNEL_ID + "/publish"

SD_CS_PIN = 5
CSV_PATH = "/sd/IoTSensorStream.csv"
SEND_INTERVAL = 10
TEST_ROWS = None

WARNING_VIBRATION = 55
CRITICAL_VIBRATION = 60

# ================= OLED =================
OLED_WIDTH = 128
OLED_HEIGHT = 64

i2c = machine.I2C(
    0,
    scl=machine.Pin(22),
    sda=machine.Pin(21)
)

oled = ssd1306.SSD1306_I2C(OLED_WIDTH, OLED_HEIGHT, i2c)


def update_oled(status, timestep, vibration, alarm_count):
    oled.fill(0)

    oled.text("IoT Monitor", 0, 0)
    oled.text("Time: {}".format(timestep), 0, 16)
    oled.text("Vib: {}".format(vibration), 0, 28)
    oled.text("Alarm: {}".format(alarm_count), 0, 40)
    oled.text(status, 0, 52)

    oled.show()


def show_startup(msg):
    oled.fill(0)
    oled.text("System Booting", 0, 0)
    oled.text(msg, 0, 20)
    oled.show()


def connect_wifi():
    print("Connecting to WiFi...", end="")
    show_startup("Connecting WiFi")

    sta_if = network.WLAN(network.STA_IF)
    sta_if.active(True)

    if not sta_if.isconnected():
        sta_if.connect(WIFI_SSID, WIFI_PASS)
        while not sta_if.isconnected():
            print(".", end="")
            time.sleep(0.5)

    print("\nWiFi connected!")
    print("IP:", sta_if.ifconfig()[0])
    
    try:
        import ntptime
        print("Setting time from NTP...", end="")
        ntptime.host = 'pool.ntp.org'
        ntptime.settime()
        print("Done")
        
        t = time.localtime()
        print("Current Time:", t)
        
    except Exception as e:
        print("Error syncing time:", e)

    show_startup("WiFi Connected")
    time.sleep(1)


def mount_sd():
    try:
        show_startup("Mounting SD")

        sd = machine.SDCard(
            slot=2,
            width=1,
            sck=machine.Pin(18),
            mosi=machine.Pin(23),
            miso=machine.Pin(19),
            cs=machine.Pin(SD_CS_PIN)
        )

        os.mount(sd, "/sd")
        print("SD card mounted successfully")
        print("Files on SD:", os.listdir("/sd"))

        show_startup("SD Mounted")
        time.sleep(1)
        return True

    except Exception as e:
        print("SD card mount failed:", e)
        show_startup("SD Failed")
        return False


def connect_mqtt():
    show_startup("Connecting MQTT")

    client = MQTTClient(
        client_id=MQTT_CLIENT_ID,
        server=MQTT_HOST,
        user=MQTT_USER,
        password=MQTT_PASS,
        keepalive=60
    )

    print("Connecting to ThingSpeak MQTT...")
    client.connect()
    print("MQTT connected!")

    show_startup("MQTT Connected")
    time.sleep(1)

    return client


def get_status(vibration):
    if vibration > CRITICAL_VIBRATION:
        return "CRITICAL"
    elif vibration > WARNING_VIBRATION:
        return "WARNING"
    else:
        return "NORMAL"


def get_alarm_code(status):
    if status == "CRITICAL":
        return 2
    elif status == "WARNING":
        return 1
    else:
        return 0

def get_current_timestamp():
    t = time.localtime()
    
    hour = t[3] + 8 
    day = t[2] 
    month = t[1]
    year = t[0]
    
    if hour >= 24:
        hour -= 24
        day += 1 
    
    return "{:04d}-{:02d}-{:02d} {:02d}:{:02d}:{:02d}".format(
        year, month, day, hour, t[4], t[5]
    )

def send_csv_data(client):
    print("Opening file:", CSV_PATH)

    alarm_count = 0

    with open(CSV_PATH, "r") as f:
        header = f.readline().strip()
        print("Header:", header)

        row_count = 0

        for line in f:
            line = line.strip()
            if not line:
                continue

            data = line.split(",")

            if len(data) < 6:
                print("Skipping invalid line:", line)
                continue

            timestep = data[0].strip()
            v_volt = data[1].strip()
            v_rotate = data[2].strip()
            v_press = data[3].strip()
            v_vibra = data[4].strip()
            label = data[5].strip()

            vibration_value = float(v_vibra)
            local_status = get_status(vibration_value)
            alarm_code = get_alarm_code(local_status)

            if alarm_code > 0:
                alarm_count += 1

            # ===== OLED实时更新 =====
            update_oled(local_status, timestep, v_vibra, alarm_count)

            payload = "field1={}&field2={}&field3={}&field4={}&field6={}".format(
                v_volt, v_rotate, v_press, v_vibra, alarm_count
            )

            try:
                client.publish(TOPIC, payload)
                row_count += 1

                timestamp = get_current_timestamp()

                print(
                    "[{}] [timestep={}] Sent | volt={} | rotate={} | pressure={} | vibration={} | alarm_count={} | status={} | label={}".format(
                        timestamp,
                        timestep,
                        v_volt,
                        v_rotate,
                        v_press,
                        v_vibra,
                        alarm_count,
                        local_status,
                        label
                    )
                )

            except Exception as e:
                print("Publish failed:", e)
                print("Reconnecting MQTT...")
                client = connect_mqtt()
                client.publish(TOPIC, payload)

            if TEST_ROWS is not None and row_count >= TEST_ROWS:
                print("Test complete. Sent {} rows.".format(row_count))
                break

            print("Waiting {} seconds...\n".format(SEND_INTERVAL))
            time.sleep(SEND_INTERVAL)

    print("Finished sending {} rows.".format(row_count))


def main():
    show_startup("Starting...")

    connect_wifi()

    sd_available = mount_sd()
    if not sd_available:
        print("Program stopped: SD card unavailable.")
        return

    client = None

    try:
        client = connect_mqtt()
        send_csv_data(client)
        print("Program finished.")

        show_startup("Finished")

    except Exception as e:
        print("Error:", e)
        show_startup("System Error")

    finally:
        if client is not None:
            try:
                client.disconnect()
                print("MQTT disconnected.")
            except:
                pass


main()