import dht
import network
import socket
import time
from machine import I2C, PWM, Pin
from umqtt_simple import MQTTClient

# --- 1. Hardware & Pin Configuration ---
dht_sensor = dht.DHT22(Pin(15))
button_warning = Pin(14, Pin.IN, Pin.PULL_DOWN)

# Servo Motor setup on GP16 (PWM @ 50Hz)
servo = PWM(Pin(16))
servo.freq(50)


def set_servo_angle(angle):
    duty = int(1000 + (angle / 180) * 8000)
    servo.duty_u16(duty)


# MPU6050 Gyroscope setup on GP4 (SDA) and GP5 (SCL)
i2c = I2C(0, sda=Pin(4), scl=Pin(5), freq=400000)
MPU6050_ADDR = 0x68

try:
    i2c.writeto_mem(MPU6050_ADDR, 0x6B, bytes([0]))  # Wake up MPU6050
    print("MPU6050 Gyroscope initialized!")
except Exception as err:
    print("MPU6050 Initialization Warning:", err)


def read_mpu6050():
    try:
        # Read 6 bytes starting from register 0x3B
        data = i2c.readfrom_mem(MPU6050_ADDR, 0x3B, 6)

        # Validate byte array length to avoid IndexError
        if len(data) < 6:
            return 0, 0, 0

        accel_x = (data[0] << 8) | data[1]
        accel_y = (data[2] << 8) | data[3]
        accel_z = (data[4] << 8) | data[5]

        if accel_x > 32767:
            accel_x -= 65536
        if accel_y > 32767:
            accel_y -= 65536
        if accel_z > 32767:
            accel_z -= 65536

        return accel_x, accel_y, accel_z
    except Exception as err:
        print("MPU6050 Read Error (Bypassed):", err)
        return 0, 0, 0  # Fallback default values


# --- 2. Wi-Fi Connection Setup ---
SSID = "Wokwi-GUEST"
PASSWORD = ""


def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    wlan.connect(SSID, PASSWORD)
    print("Connecting to Wi-Fi...", end="")
    while not wlan.isconnected():
        time.sleep(0.5)
        print(".", end="")
    print("\nConnected! Device IP:", wlan.ifconfig()[0])


connect_wifi()

# --- 3. MQTT Broker Configuration ---
MQTT_BROKER = "broker.hivemq.com"
CLIENT_ID = "Alaa_PicoW_Industrial_2026_SIC"
TOPIC_TELEMETRY = b"factory/sensor/telemetry"

mqtt_client = MQTTClient(CLIENT_ID, MQTT_BROKER)
try:
    mqtt_client.connect()
    print("Connected to MQTT Broker!")
except Exception as err:
    print("MQTT Connection Error:", err)

# --- 4. UDP Socket Configuration (Low-latency Emergency Warnings) ---
SERVER_IP = "192.168.1.15"  # Replace with your PC's local IP address
UDP_PORT = 5005
udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# --- 5. TCP Socket Configuration (Bidirectional Control Commands) ---
TCP_PORT = 8080
tcp_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

try:
    tcp_client.connect((SERVER_IP, TCP_PORT))
    tcp_client.setblocking(False)  # Non-blocking socket for real-time loop
    print("Connected to Node-RED TCP Server!")
except Exception as err:
    print("TCP Connection Warning/Error:", err)

# --- 6. Main Program Loop ---
last_telemetry_time = 0

while True:
    current_time = time.time()

    # Task A: Publish Telemetry via MQTT every 3 seconds
    if current_time - last_telemetry_time >= 3:
        try:
            dht_sensor.measure()
            temp = dht_sensor.temperature()
            hum = dht_sensor.humidity()
            ax, ay, az = read_mpu6050()

            payload = (
                '{"temp": %.1f, "humidity": %.1f, "vibration_x": %d, "vibration_y": %d}'
                % (temp, hum, ax, ay)
            )
            mqtt_client.publish(TOPIC_TELEMETRY, payload)
            print("[MQTT Sent]:", payload)

            last_telemetry_time = current_time
        except Exception as err:
            print("Telemetry Error:", err)

    # Task B: Send Immediate UDP Warning on Emergency Button Press
    if button_warning.value() == 1:
        set_servo_angle(90)  # Move motor to emergency status angle
        warning_msg = b"WARNING: Emergency Stop Button Pressed!"
        udp_socket.sendto(warning_msg, (SERVER_IP, UDP_PORT))
        print("[UDP Warning Sent]:", warning_msg)
        time.sleep(0.5)

    # Task C: Receive and Execute TCP Control Commands from Node-RED
    try:
        command = tcp_client.recv(1024).decode("utf-8").strip()
        if command:
            print("[TCP Command Received]:", command)
            if command == "START":
                set_servo_angle(180)  # Move motor to RUNNING position
                tcp_client.send(b"STARTED\n")  # Send acknowledgment back
            elif command == "STOP":
                set_servo_angle(0)  # Move motor to STOPPED position
                tcp_client.send(b"STOPPED\n")  # Send acknowledgment back
    except OSError:
        pass  # No incoming data received in this cycle

    time.sleep(0.1)