import network
import time
import socket
import struct
from machine import Pin, ADC

# ==========================================
# 1. WI-FI CONFIGURATION
# ==========================================
WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""

# ==========================================
# 2. MQTT CONFIGURATION
# ==========================================
MQTT_BROKER = "broker.hivemq.com"
MQTT_PORT = 1883

MQTT_CLIENT_ID = "EnergyMonitor_202429EE141"

TOPIC_VOLTAGE = "eszg574/202429EE141/energy/voltage"
TOPIC_CURRENT = "eszg574/202429EE141/energy/current"
TOPIC_POWER = "eszg574/202429EE141/energy/power"

PUBLISH_INTERVAL = 5

# ==========================================
# 3. ADC CONFIGURATION
# ==========================================
# Potentiometer 1: Voltage sensor, GPIO 34
# Potentiometer 2: Current sensor, GPIO 35

voltage_adc = ADC(Pin(34))
current_adc = ADC(Pin(35))

voltage_adc.atten(ADC.ATTN_11DB)
current_adc.atten(ADC.ATTN_11DB)

voltage_adc.width(ADC.WIDTH_12BIT)
current_adc.width(ADC.WIDTH_12BIT)


# ==========================================
# 4. MQTT PACKET FUNCTIONS
# ==========================================
def encode_remaining_length(length):
    result = bytearray()

    while True:
        digit = length % 128
        length //= 128

        if length > 0:
            digit |= 0x80

        result.append(digit)

        if length == 0:
            return bytes(result)


def mqtt_string(value):
    value = value.encode()
    return struct.pack("!H", len(value)) + value


def send_packet(sock, packet_type, body):
    packet = bytes([packet_type])
    packet += encode_remaining_length(len(body))
    packet += body
    sock.sendall(packet)


def read_exact(sock, count):
    result = b""

    while len(result) < count:
        chunk = sock.recv(count - len(result))

        if not chunk:
            raise OSError("MQTT connection closed")

        result += chunk

    return result


def read_packet(sock):
    first = read_exact(sock, 1)[0]

    multiplier = 1
    remaining = 0

    while True:
        digit = read_exact(sock, 1)[0]
        remaining += (digit & 127) * multiplier

        if (digit & 128) == 0:
            break

        multiplier *= 128

    body = read_exact(sock, remaining)

    return first, body


# ==========================================

# ==========================================

# ==========================================

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

    print("Connecting to Wokwi Wi-Fi...")

    if not wlan.isconnected():
        wlan.connect("Wokwi-GUEST", "")

        for attempt in range(30):
            if wlan.isconnected():
                break

            
            time.sleep(1)

    if wlan.isconnected():
        print("Wi-Fi connected!")
        print("IP address:", wlan.ifconfig()[0])
        return wlan

    raise OSError("Wi-Fi connection timeout")

# ==========================================
# 6. MQTT CONNECTION
# ==========================================
def connect_mqtt():
    print("Connecting to HiveMQ...")

    sock = socket.socket()
    sock.settimeout(10)
    sock.connect((MQTT_BROKER, MQTT_PORT))

    # MQTT 3.1.1 CONNECT packet
    variable_header = mqtt_string("MQTT")
    variable_header += bytes([4, 2])
    variable_header += struct.pack("!H", 60)

    payload = mqtt_string(MQTT_CLIENT_ID)

    send_packet(sock, 0x10, variable_header + payload)

    packet_type, body = read_packet(sock)

    if packet_type != 0x20 or len(body) < 2 or body[1] != 0:
        sock.close()
        raise OSError("MQTT connection refused")

    sock.settimeout(10)

    print("MQTT connected!")
    print("Voltage topic:", TOPIC_VOLTAGE)
    print("Current topic:", TOPIC_CURRENT)
    print("Power topic:", TOPIC_POWER)

    return sock


# ==========================================
# 7. MQTT PUBLISH FUNCTION
# ==========================================
def publish_message(sock, topic, value):
    payload = str(value).encode()
    body = mqtt_string(topic) + payload

    # PUBLISH packet, QoS 0
    send_packet(sock, 0x30, body)


# ==========================================
# 8. READ SENSORS AND CALCULATE POWER
# ==========================================
def read_energy_values():
    raw_voltage = voltage_adc.read()
    raw_current = current_adc.read()

    # Convert ADC readings (0-4095) to worksheet ranges
    voltage = (raw_voltage / 4095) * 250.0
    current = (raw_current / 4095) * 10.0

    # Round first, then calculate power from displayed values
    voltage = round(voltage, 2)
    current = round(current, 2)
    power = round(voltage * current, 2)

    return voltage, current, power


# ==========================================
# 9. MAIN MONITORING LOOP
# ==========================================
def main():
    wlan = connect_wifi()
    mqtt_sock = connect_mqtt()

    print()
    print("===================================")
    print("SMART ENERGY MONITORING STARTED")
    print("===================================")

    while True:
        try:
            # Reconnect to Wi-Fi if disconnected
            if not wlan.isconnected():
                print("Wi-Fi disconnected. Reconnecting...")
                wlan = connect_wifi()

            # Read sensors and calculate power
            voltage, current, power = read_energy_values()

            # Publish each parameter to a separate topic
            publish_message(mqtt_sock, TOPIC_VOLTAGE, voltage)
            publish_message(mqtt_sock, TOPIC_CURRENT, current)
            publish_message(mqtt_sock, TOPIC_POWER, power)

            # Display readings in Serial Monitor
            print("-----------------------------------")
            print("Voltage:", voltage, "V")
            print("Current:", current, "A")
            print("Calculated power:", power, "W")
            print("Verification V x I:",
                  round(voltage * current, 2), "W")
            print("Published voltage, current and power")
            print("Topics:")
            print(TOPIC_VOLTAGE)
            print(TOPIC_CURRENT)
            print(TOPIC_POWER)

            time.sleep(PUBLISH_INTERVAL)

        except Exception as error:
            print("Communication error:", error)
            print("Attempting to reconnect...")

            try:
                mqtt_sock.close()
            except Exception:
                pass

            time.sleep(2)

            # Retry Wi-Fi and MQTT connections
            while True:
                try:
                    wlan = connect_wifi()
                    mqtt_sock = connect_mqtt()
                    print("Reconnection successful!")
                    break

                except Exception as reconnect_error:
                    print("Reconnect failed:", reconnect_error)
                    time.sleep(3)


# ==========================================
# 10. START PROGRAM
# ==========================================
try:
    main()

except Exception as error:
    print("Program stopped:", error)