
# Script for reading data from Klipper via MQTT, storing them into CSV file and sending the file to the server
# v1, 19.4.2024, VUT FEKT Brno, Vaculík Samuel

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

import os
import sys
import csv
import threading
import time
import gzip
import shutil
import requests
import json
from requests.auth import HTTPBasicAuth
from requests.exceptions import HTTPError, ConnectionError, Timeout, RequestException
from paho.mqtt import client as mqtt_client
from queue import Queue
from jsonpath_ng import parse
from dotenv import load_dotenv


# --------------------------------------------------
#                    DEFINITIONS
# --------------------------------------------------

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

# Server url addresses (Only test server is used at this moment)
url_post = 'https://httpbin.org/post'
url_auth = 'https://httpbin.org/basic-auth/user/passwd'
SERVER_SENDING_PERIOD = 10     # In seconds

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

# Create CSV file name
file_name = sys.argv[1]      # Get the name from the main.py script
csv_file = f'/home/pi/BP_Vaculik/Logs/MQTT/Original/mqtt_{file_name}.csv'
compressed_csv_file = f'/home/pi/BP_Vaculik/Logs/MQTT/Compressed/mqtt_{file_name}.gz'

# Headers in CSV file
headers = [
           'Timestamp', 'EventTime', 'DisplayStatus_Progress', 'IdleTimeout_Printing_Time', 'PrintStats_Total_Duration',
           'PrintStats_Print_Duration', 'PrintStats_Filament_Used', 'PrintStats_Current_Layer',
           'SystemStats_System_Load', 'SystemStats_CPU_Time', 'SystemStats_Memory_Available',
           'GcodeMove_Position_X', 'GcodeMove_Position_Y', 'GcodeMove_Position_Z', 'GcodeMove_Position_E',
           'MotionReport_Live_Velocity', 'MotionReport_Live_Extruder_Velocity',
           'Extruder_Temperature', 'Extruder_Power', 'HeaterBed_Temperature', 'HeaterBed_Power',
           'RaspberryPi_Temperature', 'OctopusBoard_Temperature',
           'CoolingFan_RPM', 'CoolingFan_Speed', 'HotendFan_RPM', 'HotendFan_Speed',
           'FilamentWidth_Diameter'
          ]

parsingKey_printDuration = parse("$.status.print_stats.print_duration")

# Object for storing data
data_queue = Queue()

# Try to authenticate to the server address. The maximum number of attempts is 10.
for try_number in range(1, 11):
    try:
        response_auth = requests.get(url_auth, auth=HTTPBasicAuth(server_username, server_password))
        response_auth.raise_for_status()

        if response_auth.status_code == 200:
            print('Authentication was successful.')
            break

    except Exception as e:
        print(f'{try_number}. Error during authentication: {e}')
        if try_number < 10:
            time.sleep(1)
        else:
            print('Can not authenticate to the server! The program will be terminated.')
            exit(1)


# --------------------------------------------------
#                     FUNCTIONS
# --------------------------------------------------

# Function for sending the compressed CSV file to the server
def send_file():
    print("Sending to server...")

    try:
        with open(compressed_csv_file, 'rb') as file:
            # Prepare compressed file for sending to the server (key/name, actual file, format)
            file_for_send = {'file': (compressed_csv_file.split('/')[-1], file, 'application/gzip')}

            try:
                # Actual sending to the server
                response = requests.post(url_post, files=file_for_send, timeout=5)
                response.raise_for_status()

                print("The compressed file has been successfully sent to the server. (mqtt.py)")

            except HTTPError as e:
                print(f'HTTP error when sending a request: {e}')
            except ConnectionError as e:
                print(f'Connection error: {e}')
            except Timeout as e:
                print(f'Request timeout has expired: {e}')
            except RequestException as e:
                print(f'Error when sending a request: {e}')
            except Exception as e:
                print(f'Unexpected error: {e}')

    except FileNotFoundError as e:
        print(f'The file {compressed_csv_file} does not exist: {e}')
    except TypeError as e:
        print(f'Something is wrong with the file {compressed_csv_file}: {e}')
    except OSError as e:
        print(f'Operating system error: {e}')


# Function for compressing the CSV file
def compress_file():
    try:
        with open(csv_file, 'rb') as file_original:
            try:
                with gzip.open(compressed_csv_file, 'wb') as file_compressed:
                    # Copying data between files, data is also compressed at the same time
                    shutil.copyfileobj(file_original, file_compressed)

            except TypeError as e:
                print(f'Something is wrong with the file {compressed_csv_file}: {e}')
                pass
            except PermissionError:
                print(f"You do not have permission to write into the file {compressed_csv_file}!")
                pass
            except IsADirectoryError:
                print(f"The specified path ti the file is a directory! The path is: {compressed_csv_file}")
                pass
            except OSError as e:
                print(f"Operating system error: {e}")
                pass

    except FileNotFoundError as e:
        print(f'File {csv_file} does not exist: {e}')
    except TypeError as e:
        print(f'Something is wrong with the file {csv_file}: {e}')
        pass
    except PermissionError:
        print(f"You do not have permission to read the file {csv_file}!")
        pass
    except IsADirectoryError:
        print(f"The specified path to the file is a directory! The path is: {csv_file}")
        pass
    except OSError as e:
        print(f"Operating system error: {e}")
        pass


# Function to transfer data from queue to the CSV file
def process_data_in_queue():
    try:
        with open(csv_file, 'w', newline='') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=headers)            # Create object for writing data
            writer.writeheader()                                            # Insert headers into the CSV file
            while not data_queue.empty():                                   # Emptying the queue into the CSV file
                writer.writerow(data_queue.get())

    except TypeError as e:
        print(f'Something is wrong with the file {csv_file}: {e}')
        pass
    except ValueError:
        print(f"The file contains keys that are not in the 'fieldnames' list!")
        pass
    except PermissionError:
        print(f"You do not have permission to write into the file {csv_file}!")
        pass
    except IsADirectoryError:
        print(f"The specified path to the file is a directory! Path: {csv_file}")
        pass
    except OSError as e:
        print(f"Operating system error: {e}")
        pass


# Row creation function for CSV file
def store_data(data):

    should_create_row = True
    data = json.loads(data)

    gcode_pos = [None] * 4
    gcode_move_position = data.get('status', {}).get('gcode_move', {}).get('position', [])
    if len(gcode_move_position) > 1:
        gcode_pos = gcode_move_position

    match_print_duration = next(iter(parsingKey_printDuration.find(data)), None)
    if not match_print_duration:
        should_create_row = False

    if should_create_row:
        row = {
            'Timestamp':     int((time.time() - start_time) * 1000),           # Get the elapsed time in milliseconds
            'EventTime':    data.get('eventtime', ''),
            'DisplayStatus_Progress': data.get('status', {}).get('display_status', {}).get('progress', ''),
            'IdleTimeout_Printing_Time': data.get('status', {}).get('idle_timeout', {}).get('printing_time', ''),
            'PrintStats_Total_Duration': data.get('status', {}).get('print_stats', {}).get('total_duration', ''),
            'PrintStats_Print_Duration': data.get('status', {}).get('print_stats', {}).get('print_duration', ''),
            'PrintStats_Filament_Used': data.get('status', {}).get('print_stats', {}).get('filament_used', ''),
            'PrintStats_Current_Layer': data.get('status', {}).get('print_stats', {}).get('info.current_layer', ''),
            'SystemStats_System_Load': data.get('status', {}).get('system_stats', {}).get('sysload', ''),
            'SystemStats_CPU_Time': data.get('status', {}).get('system_stats', {}).get('cputime', ''),
            'SystemStats_Memory_Available': data.get('status', {}).get('system_stats', {}).get('memavail', ''),
            'GcodeMove_Position_X': gcode_pos[0],
            'GcodeMove_Position_Y': gcode_pos[1],
            'GcodeMove_Position_Z': gcode_pos[2],
            'GcodeMove_Position_E': gcode_pos[3],
            'MotionReport_Live_Velocity': data.get('status', {}).get('motion_report', {}).get('live_velocity', ''),
            'MotionReport_Live_Extruder_Velocity': data.get('status', {}).get('motion_report', {}).get('live_extruder_velocity', ''),
            'Extruder_Temperature': data.get('status', {}).get('extruder', {}).get('temperature', ''),
            'Extruder_Power': data.get('status', {}).get('extruder', {}).get('power', ''),
            'HeaterBed_Temperature': data.get('status', {}).get('heater_bed', {}).get('temperature', ''),
            'HeaterBed_Power': data.get('status', {}).get('heater_bed', {}).get('power', ''),
            'RaspberryPi_Temperature': data.get('status', {}).get('temperature_sensor raspberry_pi', {}).get('temperature', ''),
            'OctopusBoard_Temperature':  data.get('status', {}).get('temperature_sensor mcu_temp', {}).get('temperature', ''),
            'CoolingFan_RPM': data.get('status', {}).get('fan', {}).get('rpm', ''),
            'CoolingFan_Speed': data.get('status', {}).get('fan', {}).get('speed', ''),
            'HotendFan_RPM': data.get('status', {}).get('hotend_fan', {}).get('rpm', ''),
            'HotendFan_Speed': data.get('status', {}).get('hotend_fan', {}).get('speed', ''),
            'FilamentWidth_Diameter': data.get('status', {}).get('hall_filament_width_sensor', {}).get('Diameter', '')
        }

        data_queue.put(row)                                                 # Insert the created row into the queue


# --------------------------------------------------
#                THREAD FUNCTION
# --------------------------------------------------

def periodic_sending_to_the_server():
    while True:
        time.sleep(SERVER_SENDING_PERIOD)

        if not data_queue.empty():
            process_data_in_queue()
            compress_file()
            send_file()
        else:
            print('Queue is empty.')


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

def connect_mqtt() -> mqtt_client:
    def on_connect(client_mqtt, userdata, flags, rc):
        if rc == 0:
            print("Connected to MQTT Broker!")
        else:
            print("Failed to connect! %d\n", rc)

    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):
        store_data(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()

    # Sets up an independent loop for sending data to the server
    threading.Thread(target=periodic_sending_to_the_server, daemon=True).start()

    # Actual start of the script
    start_mqtt()
