import network
import time
import socket
import ujson
from umqtt.simple import MQTTClient

# ---------------- Wi-Fi SETTINGS ----------------
WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""

# ---------------- MQTT SETTINGS ----------------
MQTT_BROKER = "broker.hivemq.com"
MQTT_PORT = 1883
MQTT_CLIENT_ID = "bridge_coordinator_2026"

# Topics for Node 1
TOPIC_N1_VIB = b"bridge/node1/vibration"
TOPIC_N1_TEMP = b"bridge/node1/temperature"
TOPIC_N1_STRAIN = b"bridge/node1/strain"

# Topics for Node 2
TOPIC_N2_VIB = b"bridge/node2/vibration"
TOPIC_N2_TEMP = b"bridge/node2/temperature"
TOPIC_N2_STRAIN = b"bridge/node2/strain"

# ---------------- THINGSPEAK SETTINGS ----------------
THINGSPEAK_WRITE_API_KEY = "01PXOE17PVANL2UZ"
THINGSPEAK_HOST = "api.thingspeak.com"
UPLOAD_INTERVAL = 20

# ---------------- ALERT THRESHOLDS ----------------
VIBRATION_THRESHOLD = 70
TEMPERATURE_THRESHOLD = 50
STRAIN_THRESHOLD = 600

# ---------------- STORAGE FOR BOTH NODES ----------------
node1_data = {"vibration": None, "temperature": None, "strain": None}
node2_data = {"vibration": None, "temperature": None, "strain": None}

last_upload_time = 0

# ---------------- WIFI CONNECTION ----------------
def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)

    if wlan.isconnected():
        print("Wi-Fi already connected:", wlan.ifconfig())
        return wlan

    print("Connecting to Wi-Fi...")
    wlan.connect(WIFI_SSID, WIFI_PASSWORD)

    timeout = 20
    start_time = time.time()

    while not wlan.isconnected():
        if time.time() - start_time > timeout:
            raise Exception("Wi-Fi connection timeout")
        time.sleep(1)

    print("Wi-Fi Connected!")
    print("IP Address:", wlan.ifconfig()[0])
    return wlan

# ---------------- MQTT CALLBACK ----------------
def mqtt_callback(topic, message):
    global node1_data, node2_data

    try:
        value = float(message)

        # Node 1 updates
        if topic == TOPIC_N1_VIB:
            node1_data["vibration"] = value
            print("Node 1 Vibration:", value)
        elif topic == TOPIC_N1_TEMP:
            node1_data["temperature"] = value
            print("Node 1 Temperature:", value)
        elif topic == TOPIC_N1_STRAIN:
            node1_data["strain"] = value
            print("Node 1 Strain:", value)

        # Node 2 updates
        elif topic == TOPIC_N2_VIB:
            node2_data["vibration"] = value
            print("Node 2 Vibration:", value)
        elif topic == TOPIC_N2_TEMP:
            node2_data["temperature"] = value
            print("Node 2 Temperature:", value)
        elif topic == TOPIC_N2_STRAIN:
            node2_data["strain"] = value
            print("Node 2 Strain:", value)

    except Exception as error:
        print("Invalid message received:", error)

# ---------------- MQTT CONNECTION ----------------
def connect_mqtt():
    client = MQTTClient(
        MQTT_CLIENT_ID,
        MQTT_BROKER,
        port=MQTT_PORT,
        keepalive=60
    )

    client.set_callback(mqtt_callback)
    client.connect()
    print("MQTT Connected Successfully!")

    # Subscribe to all 6 topics (3 from Node 1, 3 from Node 2)
    for topic in [TOPIC_N1_VIB, TOPIC_N1_TEMP, TOPIC_N1_STRAIN, 
                  TOPIC_N2_VIB, TOPIC_N2_TEMP, TOPIC_N2_STRAIN]:
        client.subscribe(topic)
        print("Subscribed to:", topic)

    return client

# ---------------- ALERT CHECK ----------------
def check_alerts():
    alert = 0
    
    for name, data in [("Node 1", node1_data), ("Node 2", node2_data)]:
        if data["vibration"] and data["vibration"] > VIBRATION_THRESHOLD:
            alert = 1
            print(f"ALERT: High vibration on {name}!")
        if data["temperature"] and data["temperature"] > TEMPERATURE_THRESHOLD:
            alert = 1
            print(f"ALERT: High temperature on {name}!")
        if data["strain"] and data["strain"] > STRAIN_THRESHOLD:
            alert = 1
            print(f"ALERT: High strain on {name}!")

    if alert == 0:
        print("Bridge Status: NORMAL across all nodes")

    return alert

# ---------------- THINGSPEAK UPLOAD ----------------
def upload_to_thingspeak(alert):
    # Using Node 1 values as primary fields, or you can adjust field mapping
    if None in node1_data.values():
        print("Waiting for complete Node 1 telemetry before uploading.")
        return

    try:
        data = (
            "api_key=" + THINGSPEAK_WRITE_API_KEY
            + "&field1=" + str(node1_data["vibration"])
            + "&field2=" + str(node1_data["temperature"])
            + "&field3=" + str(node1_data["strain"])
            + "&field4=" + str(alert)
        )

        address = socket.getaddrinfo(THINGSPEAK_HOST, 80)[0][-1]
        sock = socket.socket()
        sock.settimeout(10)
        sock.connect(address)

        request = (
            "POST /update HTTP/1.1\r\n"
            "Host: " + THINGSPEAK_HOST + "\r\n"
            "Connection: close\r\n"
            "Content-Type: application/x-www-form-urlencoded\r\n"
            "Content-Length: " + str(len(data)) + "\r\n"
            "\r\n"
            + data
        )

        sock.send(request.encode())
        response = sock.recv(512)
        sock.close()

        if b"200 OK" in response:
            print("ThingSpeak upload successful.")
        else:
            print("ThingSpeak response was not HTTP 200.")

    except Exception as error:
        print("ThingSpeak upload error:", error)

# ---------------- MAIN PROGRAM ----------------
print("SMART BRIDGE MONITORING SYSTEM - MULTI-NODE COORDINATOR")

while True:
    try:
        connect_wifi()
        mqtt_client = connect_mqtt()

        print("Coordinator listening...")

        while True:
            mqtt_client.check_msg()
            current_time = time.time()

            if current_time - last_upload_time >= UPLOAD_INTERVAL:
                alert_status = check_alerts()
                upload_to_thingspeak(alert_status)
                last_upload_time = current_time

            time.sleep(1)

    except Exception as error:
        print("Coordinator error:", error)
        print("Reconnecting in 5 seconds...")
        time.sleep(5)