import network
import rp2
import time
import os
import json
import machine
from machine import Pin
from umqtt.robust import MQTTClient
from pico_funcs import read_cpu_temp, wlan_up, led_error_code
from hcsr04_funcs import read_hc_sr04
from dotenv import load_dotenv

load_dotenv()

# Change this GPIO PIN where your DHT22 sensor is connected
TRIG_PIN = 3
ECHO_PIN = 2

# Debug Mode
DEBUG = False


def main():

    # Start Up Activities
    if DEBUG:
        print("Start up")
    count = 0
    led = machine.Pin("LED", machine.Pin.OUT)
    led.value(False)
    led_error_code(led, 1)

    # Set Wi-Fi Country and Create a Wireless Interface
    rp2.country("US")
    wlan = network.WLAN(network.STA_IF)

    # Create MQTT Client
    mqtt_client = MQTTClient(
        client_id=os.getenv("MQTT_CLIENT_ID"),
        server=os.getenv("MQTT_HOST_NAME"),
    )

    # Create Home Assistant MQTT Client
    ha_mqtt_client = MQTTClient(
        client_id=os.getenv("HA_MQTT_CLIENT_ID"),
        server=os.getenv("HA_MQTT_HOST_NAME"),
        keepalive=os.getenv("HA_MQTT_KEEP"),
        user=os.getenv("HA_MQTT_USERNAME"),
        password=os.getenv("HA_MQTT_PASSWORD"),
    )

    # Create HC-SR04
    trig = Pin(TRIG_PIN, Pin.OUT)
    echo = Pin(ECHO_PIN, Pin.IN, Pin.PULL_DOWN)

    # Let's Go!
    if DEBUG:
        print("Enter main loop")

    # First Time flag allows the 1st order filter to be initialized.
    first_time = True
    FILTER_CONSTANT = 0.9

    while True:

        # Loop Clean-Up and Prep
        led.value(False)
        count += 1
        if DEBUG:
            print(f"\nStarting loop #{count}... \nWiFI Status is {wlan.status()}.")

        # Create Local Flags for Control
        wifi_ready = False
        mqtt_ready = False
        hc_sr04_ready = False

        # CPU Reading.
        try:
            cpu_temp = read_cpu_temp()
        except:
            continue

        # HC-SR04 Reading.
        try:
            sensor_value = read_hc_sr04(trig, echo)
            if not first_time:
                distance = (
                    distance * FILTER_CONSTANT + (1 - FILTER_CONSTANT) * sensor_value
                )
            else:
                distance = sensor_value
                first_time = False
            hc_sr04_ready = True
        except:
            continue

        # WiFi Connection.
        try:
            if wlan.status() != 3:
                ifconfig = wlan_up(wlan)
                if ifconfig is None:
                    if DEBUG:
                        print("Trouble to connecting WiFi: {}".format(e))
                    led_error_code(led, 3)
                    time.sleep(10)
                    continue
                else:
                    wifi_ready = True
                    if DEBUG:
                        print("Connected to WiFi: {}".format(ifconfig))
            else:
                wifi_ready = True
                ifconfig = wlan.ifconfig()
                if DEBUG:
                    print(
                        f"Skipping WiFi Activation since WiFi Status is {wlan.status()} \n{ifconfig}."
                    )
        except:
            continue

        # MQTT Connection to Primary Broker.
        try:
            mqtt_client.connect(clean_session=False)
            mqtt_ready = True
            mqtt = "mosquitto"
            if DEBUG:
                print(f"Connected to Mosquitto MQTT broker.")

        except Exception as e:
            if DEBUG:
                print("Trouble to connecting to Mosquitto MQTT: {}".format(e))

        # MQTT Connection to Back Up Broker.
        if not mqtt_ready:
            try:
                try:
                    ha_mqtt_client.connect(clean_session=False)
                    mqtt_ready = True
                    mqtt = "ha"
                    if DEBUG:
                        print(f"Connected to the back up, Home Assistant, MQTT broker.")

                except Exception as f:
                    if DEBUG:
                        print(
                            "Trouble to connecting to Home Assistant MQTT: {}".format(f)
                        )
                    led_error_code(led, 2)
                    time.sleep(5)
                    continue  # Start back at the top of the While Loop

            except:
                continue

        # Ready to Publish?
        try:
            if wifi_ready and mqtt_ready and hc_sr04_ready:
                # Timestamp
                time_now = time.localtime()
                timestamp = "{}/{}/{} {}:{}:{}".format(
                    time_now[1],
                    time_now[2],
                    time_now[0],
                    time_now[3],
                    time_now[4],
                    time_now[5],
                )

                # Build JSON Payloads
                hc_sr04_data = {
                    "Location": os.getenv("MQTT_ZONE_ID"),
                    "Distance": distance,
                }
                hc_sr04_payload = json.dumps(hc_sr04_data)

                hw_data = {
                    "Timestamp": timestamp,
                    "CPU Temperature": cpu_temp,
                    "Device": os.getenv("MQTT_HW_ID"),
                    "WiFi Information": ifconfig,
                }
                hw_payload = json.dumps(hw_data)

                if DEBUG:
                    print(f"Trying to publish...")

                if mqtt == "mosquitto":
                    # Publish
                    mqtt_client.publish(
                        "home/{}".format(os.getenv("MQTT_ZONE_ID")),
                        hc_sr04_payload,
                        retain=True,
                    )
                    mqtt_client.publish(
                        "hw/{}".format(os.getenv("MQTT_HW_ID")), hw_payload, retain=True
                    )
                    mqtt_client.disconnect()
                    if DEBUG:
                        print(f"MQTT Disconnected.")

                if mqtt == "ha":
                    # Publish
                    ha_mqtt_client.publish(
                        "home/{}".format(os.getenv("HA_MQTT_ZONE_ID")),
                        hc_sr04_payload,
                        retain=True,
                    )
                    ha_mqtt_client.publish(
                        "hw/{}".format(os.getenv("HA_MQTT_HW_ID")),
                        hw_payload,
                        retain=True,
                    )
                    ha_mqtt_client.disconnect()
                    if DEBUG:
                        print("MQTT Disconnected.")

        except:
            continue

        # Sleep for a bit.
        if DEBUG:
            print(f"Finished loop #{count}.")
        time.sleep(5)


main()