import network
import rp2
import time
import json
import machine
from machine import Pin
from umqtt.robust import MQTTClient
import wifi_config
import mqtt_config
import ha_mqtt_config
from pico_funcs import read_cpu_temp, wlan_up, led_error_code
from hcsr04_funcs import read_hc_sr04

# 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 = mqtt_config.MQTT_CLIENT_ID,
        server = mqtt_config.MQTT_HOST_NAME,
        )

    # Create Home Assistant MQTT Client
    ha_mqtt_client = MQTTClient(
        client_id = ha_mqtt_config.MQTT_CLIENT_ID,
        server = ha_mqtt_config.MQTT_HOST_NAME,
        keepalive = ha_mqtt_config.MQTT_KEEP,
        user = ha_mqtt_config.MQTT_USERNAME,
        password = ha_mqtt_config.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':mqtt_config.MQTT_ZONE_ID,
                    'Distance':distance,
                }
                hc_sr04_payload = json.dumps(hc_sr04_data)
                
                hw_data = {
                    'Timestamp':timestamp,
                    'CPU Temperature':cpu_temp,
                    'Device':mqtt_config.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(mqtt_config.MQTT_ZONE_ID),hc_sr04_payload, retain=True)
                    mqtt_client.publish("hw/{}".format(mqtt_config.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(mqtt_config.MQTT_ZONE_ID),hc_sr04_payload, retain=True)
                    ha_mqtt_client.publish("hw/{}".format(mqtt_config.MQTT_HW_ID),hw_payload, retain=True)            
                    ha_mqtt_client.disconnect()
                    if DEBUG:
                        print(f'MQTT Disconnected.')



        except:
            continue

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

main()