#!/usr/bin/env python3
import subprocess
import json
import sys
import paho.mqtt.client as mqtt
import time

def check_backup_vg():
    result = subprocess.run(["vgdisplay", "-cA"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
    vgdisplay_output = result.stdout.strip().split("\n")
    return any(line.split(":")[0].strip() == "backup" for line in vgdisplay_output)

def send_mqtt_message(topic, payload):
    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)  # Use the latest API version
    client.connect("mqtt.chaos", 1883, 60)
    client.publish(topic, payload)
    client.disconnect()

def get_args(name='default', first='on'):
    return first

mode = get_args(*sys.argv)
print(f"Mode is: {mode}")
if mode not in ["on", "off"]:
    print("Error: mode must be either 'on' or 'off'")
    sys.exit(1)

# Run the systemctl command and capture the output
try:
    result = subprocess.run(
        ["systemctl", "list-units", "--type=automount", "--no-legend", "--output=json"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        check=True
    )
except subprocess.CalledProcessError as e:
    print(f"Error executing command: {e.stderr}")
    exit(1)

# Parse the JSON output
try:
    units = json.loads(result.stdout)
except json.JSONDecodeError as e:
    print(f"Error parsing JSON: {e}")
    exit(1)

# Filter units that are loaded, active, and sub: running, and start with 'backup-'
running_units = []
for unit in units:
    if (
        unit.get("load") == "loaded" and
        unit.get("active") == "active" and
        unit.get("sub") == "running" and
        unit.get("unit", "").startswith("backup-")
    ):
        running_units.append(unit)

# Output the results
print("Units with load: loaded, active: active, sub: running:")
for unit in running_units:
    print(f"  - {unit['unit']}: {unit['description']}")

# Send MQTT message based on mode
if mode == "on":
    print(f"Turning backup Disk ON")
    send_mqtt_message("switch_backup/switch/switch_backup_power/command", "on")
    time.sleep(15)
    if not check_backup_vg():
        try:
            print(f"Activating VG backup")
            subprocess.run(["vgchange", "-ay", "backup"], check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error executing vgchange: {e.stderr}")
            exit(1)
else:
    # OFF case: Only execute if nothing is mounted
    if not running_units and check_backup_vg():
        print(f"Turning backup Disk OFF")
        subprocess.run(["sync"], check=True)
        try:
            subprocess.run(["vgchange", "-an", "backup"], check=True)
            subprocess.run(["sync"], check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error executing vgchange: {e.stderr}")
            exit(1)
        send_mqtt_message("switch_backup/switch/switch_backup_power/command", "off")