#!/usr/bin/env python3

import os
import sys
import yaml
import paho.mqtt.client as mqtt
import time

def load_config():
    """Load configuration from YAML file."""
    try:
        with open('config.yaml', 'r') as f:
            return yaml.safe_load(f)
    except Exception as e:
        print(f"Error loading config: {e}")
        sys.exit(1)

def on_connect(client, userdata, flags, rc):
    """Callback for when the client connects to the MQTT broker."""
    if rc == 0:
        print("Connected to MQTT broker")
    else:
        print(f"Failed to connect, return code {rc}")
        sys.exit(1)

def on_message(client, userdata, message):
    """Callback for when a message is received."""
    userdata.append(message.topic)

def get_topics(client, timeout=1):
    """Get all system2mqtt topics from the MQTT broker."""
    topics = []
    client.user_data_set(topics)
    client.on_message = on_message
    
    # Subscribe to all topics
    client.subscribe('#')
    client.loop_start()
    
    # Wait for messages
    time.sleep(timeout)
    
    client.loop_stop()
    client.unsubscribe('#')
    
    # Filter for system2mqtt topics
    return [topic for topic in topics if 'system2mqtt' in topic]

def cleanup_topics(config):
    """Clean up MQTT topics."""
    # Create MQTT client
    client = mqtt.Client()
    client.username_pw_set(config['mqtt']['username'], config['mqtt']['password'])
    client.on_connect = on_connect

    # Connect to MQTT broker
    try:
        print(f"Connecting to MQTT broker at {config['mqtt']['host']}:{config['mqtt']['port']}...")
        client.connect(config['mqtt']['host'], config['mqtt']['port'], 60)
    except Exception as e:
        print(f"Failed to connect to MQTT broker: {e}")
        sys.exit(1)

    # Get all system2mqtt topics
    print("Finding system2mqtt topics...")
    topics = get_topics(client)
    
    if not topics:
        print("No system2mqtt topics found")
        return

    # Delete each topic
    for topic in topics:
        print(f"Deleting topic: {topic}")
        client.publish(topic, "", retain=True)
        time.sleep(0.1)  # Small delay between deletions

    # Wait for messages to be published
    time.sleep(1)

    # Stop the MQTT client loop and disconnect
    client.disconnect()

    print("Cleanup completed")

if __name__ == "__main__":
    config = load_config()
    cleanup_topics(config)