# mqtt_motor_control.py
import paho.mqtt.client as mqtt
import time

# Настройки MQTT
MQTT_BROKER = "192.168.31.225"  # Адрес MQTT брокера
MQTT_PORT = 1883                 # Порт MQTT
MQTT_USER = "test"            # Имя пользователя
MQTT_PASSWORD = "1234"         # Пароль

# Топики
MQTT_DRIVER_TOPIC = "motor/control/driver"             # Топик для управления драйвером
MQTT_RPM_TOPIC = "motor/control/rpm"                   # Топик для установки скорости оборотов в минуту
MQTT_RESET_STEPS_TOPIC = "motor/control/totalsteps/reset"  # Топик для сброса шагов

# Топики для управления TMC2209
MQTT_TMC_CURRENT_TOPIC = "motor/control/tmc/current_percent"         # Топик для установки тока TMC2209
MQTT_TMC_MICROSTEPS_TOPIC = "motor/control/tmc/microsteps"  # Топик для установки микрошагов TMC2209
MQTT_TMC_STALLGUARD_TOPIC = "motor/control/tmc/stallguard" # Топик для установки StallGuard TMC2209
MQTT_TMC_ENABLE_TOPIC = "motor/control/tmc/enable"
MQTT_TMC_STEALTHCHOP_TOPIC = "motor/control/tmc/stealthchop"
MQTT_TMC_COOLSTEP_TOPIC = "motor/control/tmc/coolstep"

def on_connect(client, userdata, flags, rc):
    if rc == 0:
        print("Connected to MQTT broker")
        # Подписка на все топики для получения обратной связи
        client.subscribe("motor/feedback/#")
    else:
        print(f"Failed to connect, return code {rc}")

def send_command(topic, command):
    """Отправка команды в MQTT топик"""
    client.publish(topic, command)
    print(f"Sent command: {command} to topic {topic}")

def handle_command(command):
    """Обработка команд"""
    parts = command.split()
    if not parts:
        return
    
    cmd = parts[0].lower()
    
    # Используем словарь для обработки команд
    command_handlers = {
        "turnon": lambda: send_command(MQTT_DRIVER_TOPIC, "turnon"),
        "turnoff": lambda: send_command(MQTT_DRIVER_TOPIC, "turnoff"),
        "resetsteps": lambda: send_command(MQTT_RESET_STEPS_TOPIC, "reset"),
    }
    
    # Обработка команд с параметрами
    if cmd == "rpm" and len(parts) > 1:
        try:
            rpm = int(parts[1])
            send_command(MQTT_RPM_TOPIC, str(rpm))
        except ValueError:
            print("Invalid RPM value. Please enter a number.")
    elif cmd == "current" and len(parts) > 1:
        try:
            current = int(parts[1])
            send_command(MQTT_TMC_CURRENT_TOPIC, str(current))
        except ValueError:
            print("Invalid current value. Please enter a number.")
    elif cmd == "microsteps" and len(parts) > 1:
        try:
            microsteps = int(parts[1])
            send_command(MQTT_TMC_MICROSTEPS_TOPIC, str(microsteps))
        except ValueError:
            print("Invalid microsteps value. Please enter a number.")
    elif cmd == "stallguard" and len(parts) > 1:
        try:
            stallguard = int(parts[1])
            send_command(MQTT_TMC_STALLGUARD_TOPIC, str(stallguard))
        except ValueError:
            print("Invalid stallguard value. Please enter a number.")
    elif cmd == "enable" and len(parts) > 1:
        try:
            enable = parts[1].lower()
            if enable in ["on", "off"]:
                send_command(MQTT_TMC_ENABLE_TOPIC, enable)
            else:
                print("Invalid enable value. Please use 'on' or 'off'.")
        except ValueError:
            print("Invalid enable value. Please use 'on' or 'off'.")
    elif cmd == "stealthchop" and len(parts) > 1:
        try:
            stealthchop = parts[1].lower()
            if stealthchop in ["on", "off"]:
                send_command(MQTT_TMC_STEALTHCHOP_TOPIC, stealthchop)
            else:
                print("Invalid stealthchop value. Please use 'on' or 'off'.")
        except ValueError:
            print("Invalid stealthchop value. Please use 'on' or 'off'.")
    elif cmd == "coolstep" and len(parts) > 1:
        try:
            coolstep = parts[1].lower()
            if coolstep in ["on", "off"]:
                send_command(MQTT_TMC_COOLSTEP_TOPIC, coolstep)
            else:
                print("Invalid coolstep value. Please use 'on' or 'off'.")
        except ValueError:
            print("Invalid coolstep value. Please use 'on' or 'off'.")
    elif cmd in command_handlers:
        command_handlers[cmd]()
    else:
        print("Unknown command. Use 'start', 'stop', 'rpm <value>', 'speed <value>', 'resetsteps', 'current <value>', 'microsteps <value>', 'stallguard <value>', 'enable <on/off>', 'stealthchop <on/off>', 'coolstep <on/off>', or 'exit'")

def main():
    # Создание клиента MQTT
    global client
    client = mqtt.Client()
    
    # Установка учетных данных
    client.username_pw_set(MQTT_USER, MQTT_PASSWORD)
    
    # Установка функции обратного вызова при подключении
    client.on_connect = on_connect
    
    # Подключение к брокеру
    try:
        client.connect(MQTT_BROKER, MQTT_PORT, 60)
        client.loop_start()  # Запуск цикла обработки в отдельном потоке
        
        print("MQTT Motor Controller")
        print("Commands:")
        print("  turnon - включить управление двигателем")
        print("  turnoff - выключить управление двигателем")
        print("  rpm <value> - установить скорость оборотов (например: rpm 200)")
        print("  resetsteps - сбросить счетчик шагов")
        print("  current <value> - установить ток в % TMC2209 (например: current 50)")
        print("  microsteps <value> - установить микрошаги TMC2209 (например: microsteps 16)")
        print("  stallguard <value> - установить StallGuard TMC2209 (например: stallguard 5)")
        print("  enable <on/off> - включить/выключить TMC2209")
        print("  stealthchop <on/off> - включить/выключить stealthchop")
        print("  coolstep <on/off> - включить/выключить coolstep")
        print("  exit - выход")
        
        while True:
            command = input("\nEnter command: ").strip()
            
            if command == "exit":
                break
                
            handle_command(command)
                
    except KeyboardInterrupt:
        print("\nShutting down...")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        client.loop_stop()
        client.disconnect()
        print("Disconnected")

if __name__ == "__main__":
    main()