
# Main script for starting and ending measurement
# v1, 19.4.2024, VUT FEKT Brno, Vaculík Samuel


# --------------------------------------------------
#                     LIBRARIES
# --------------------------------------------------

import os
import json
import subprocess
import threading
import time
from jsonpath_ng import parse
from datetime import datetime
from dotenv import load_dotenv
from paho.mqtt import client as mqtt_client


# --------------------------------------------------
#                     VARIABLES
# --------------------------------------------------

# MQTT parameters
broker = 'octopus.local'
port = 1883
topic = "octopus/klipper/status"

# Get login data to the mosquitto mqtt broker from the '.env' file
dotenv_path = '/home/pi/BP_Vaculik/.env'
load_dotenv(dotenv_path)
mosquitto_username = os.getenv('MOSQUITTO_USERNAME')
mosquitto_password = os.getenv('MOSQUITTO_PASSWORD')

# Print info variables
parsingKey_printDuration = parse("$.status.print_stats.print_duration")
parsingKey_totalDuration = parse("$.status.print_stats.total_duration")
parsingKey_filename = parse("$.status.print_stats.filename")

# Auxiliary variables
csv_filename_exists = None
csv_filename = None
printing = False

# List for Arduino, Microphones, MQTT subprocesses
processes = []


# --------------------------------------------------
#               SUBPROCESS FUNCTIONS
# --------------------------------------------------

def processes_start(name):
    print("Starting arduino.py, microphones.py and mqtt.py")
    global processes

    processes.append(subprocess.Popen(['python3', '/home/pi/BP_Vaculik/Python/Arduino/arduino.py', name]))
    processes.append(subprocess.Popen(['python3', '/home/pi/BP_Vaculik/Python/MQTT/mqtt.py', name]))
    processes.append(subprocess.Popen(['python3', '/home/pi/BP_Vaculik/Python/Microphones/microphones.py', name]))


def process_stop(process, timeout):
    print(f'Closing process: {process}')
    process.terminate()                 # Try to cancel/close the process

    try:
        process.wait(timeout=timeout)
    except subprocess.TimeoutExpired:   # Process has not been terminated properly and will be terminated by force
        process.kill()
        process.wait()                  # Wait for the actual end of the process after the kill


# --------------------------------------------------
#                     MQTT
# --------------------------------------------------

# Function for checking if message contains parameters to start arduino.py, microphones.py and mqtt.py scripts
def check_msg(message):
    global printing, csv_filename, csv_filename_exists

    try:
        data = json.loads(message)

        # Check if message contains "filename" to create name for CSV files
        if not csv_filename_exists:
            match_filename = next(iter(parsingKey_filename.find(data)), None)

            if match_filename and match_filename.value != '':
                csv_filename = match_filename.value
                csv_filename = csv_filename.rsplit('.', 1)[0]       # Delete ".gcode" extension

                csv_filename_exists = True

        # Check if message contains info about print duration (if yes, print has started)
        match_print_duration = next(iter(parsingKey_printDuration.find(data)), None)

        if match_print_duration and match_print_duration.value != 0:
            if not printing:
                if not csv_filename:  # If CSV file name is still "None", create secondary name
                    print(f'Filename was not found! Only unique ID will be used.')
                    csv_filename = "Untitled"

                id_number = datetime.now().strftime('%d%m%H%M')                  # Create a unique number for each print
                if '_id' in csv_filename:
                    csv_filename = csv_filename.rsplit("_", maxsplit=1)[0]
                csv_filename = '{}_id{}'.format(csv_filename, id_number)    # Add a unique number to the file name
                print(f'CSV file name: {csv_filename}')

                processes_start(csv_filename)                                   # Starts processes
                printing = True

        else:
            if printing:
                match_total_duration = next(iter(parsingKey_totalDuration.find(data)), None)

                if not match_total_duration:    # End processes, when "total duration" is missing (= end of the print)
                    for process in processes:
                        threading.Thread(target=process_stop, args=(process, 5)).start()

                    processes[:] = []                  # Reset List for another print

                    printing = False
                    csv_filename_exists = False        # Reset name for another print

    except json.decoder.JSONDecodeError as e:
        print(f'Error when decoding JSON message: {e}')
        pass


def connect_mqtt():
    def on_connect(client_mqtt, userdata, flags, rc):
        if rc == 0:
            print('Connected to MQTT Broker!')
            pass
        else:
            print(f'Failed to connect! {rc}')
            pass

    client = mqtt_client.Client()
    client.username_pw_set(mosquitto_username, mosquitto_password)
    client.on_connect = on_connect
    client.connect(broker, port)
    return client


def subscribe(client: mqtt_client):
    def on_message(client_mqtt, userdata, msg):
        check_msg(msg.payload.decode())

    client.subscribe(topic)
    client.on_message = on_message


def start_mqtt():
    client = connect_mqtt()
    subscribe(client)
    client.loop_forever()


# --------------------------------------------------
#                     STARTUP
# --------------------------------------------------

if __name__ == '__main__':
    # Store program start time for timestamps
    start_time = time.time()

    # Actual start of the script
    start_mqtt()

