#!/usr/bin/env python3
import canopen
import time
import argparse
from enum import Enum

# --- Constants & Enums (Following torso_platform.py style) ---

class EZMotionCommands(Enum):
    """Object Dictionary mapping for EZMotion C2 Series."""
    CONTROL_WORD = 0x6040
    INTERNAL_STATUS = 0x200B  # Sub-indices: 0x09 Max V, 0x0A Min V, 0x0B Faults
    NODE_ID_SETTING = 0x2101
    STORE_PARAMETERS = 0x200D

class EZMotionValues:
    """Special hex codes for specific commands."""
    SAVE_TO_NVM = 0x65766173  # "save" in hex
    FAULT_RESET = 0x80        # Bit 7 high
    
# Configuration Defaults
DEFAULT_EDS_PATH = "./MMS760400-48-C2-1.eds"
BITRATE = 1000000
VOLTAGE_LOWER_LIMIT = 52
VOLTAGE_UPPER_LIMIT = 65

class EZMotionConfigurator:
    def __init__(self, channel='can3', bitrate=BITRATE, eds_path=DEFAULT_EDS_PATH):
        self.network = canopen.Network()
        self.channel = channel
        self.bitrate = bitrate
        self.eds_path = eds_path

    def connect(self):
        print(f"Connecting to {self.channel} at {self.bitrate}bps...")
        self.network.connect(channel=self.channel, bustype="socketcan", bitrate=self.bitrate)

    def disconnect(self):
        self.network.disconnect()
        print("Network disconnected.")

    def update_node_id(self, current_id, new_id):
        try:
            # 1. Add the node with the current ID
            node = self.network.add_node(current_id, self.eds_path)
            print(f"Successfully added Node {current_id}")

            # 2. Read current diagnostics
            fault_register = node.sdo[EZMotionCommands.INTERNAL_STATUS.value][0x0B].raw
            current_node_id = node.sdo[EZMotionCommands.NODE_ID_SETTING.value].raw
            max_v = node.sdo[EZMotionCommands.INTERNAL_STATUS.value][0x09].raw
            min_v = node.sdo[EZMotionCommands.INTERNAL_STATUS.value][0x0A].raw

            print(f"--- Diagnostic Report (Node {current_id}) ---")
            print(f"Fault Register: {hex(fault_register)}")
            print(f"Current Configured ID: {current_node_id}")
            print(f"Voltage Protection: {min_v}V (Min) - {max_v}V (Max)")

            # 3. Update Voltage Protection Limits
            print(f"\nUpdating voltage limits to {VOLTAGE_LOWER_LIMIT}V - {VOLTAGE_UPPER_LIMIT}V...")
            node.sdo[EZMotionCommands.INTERNAL_STATUS.value][0x0A].raw = VOLTAGE_LOWER_LIMIT
            node.sdo[EZMotionCommands.INTERNAL_STATUS.value][0x09].raw = VOLTAGE_UPPER_LIMIT

            # 4. Change the Node ID
            print(f"Changing Node ID from {current_id} to {new_id}...")
            node.sdo[EZMotionCommands.NODE_ID_SETTING.value].raw = new_id

            # 5. Save to NVM (Non-Volatile Memory)
            # This is critical so the ID doesn't revert to default on power cycle
            print("Saving changes to NVM...")
            node.sdo[EZMotionCommands.STORE_PARAMETERS.value].raw = EZMotionValues.SAVE_TO_NVM
            
            # 6. Reset Faults (Optional, but included from your original logic)
            # node.sdo[EZMotionCommands.CONTROL_WORD.value].raw = EZMotionValues.FAULT_RESET
            
            # Short sleep to allow the motor to finish the NVM write
            time.sleep(0.5)
            print(f"Success. Motor will respond as ID {new_id} after next power cycle or reset.")

        except Exception as e:
            print(f"Error communicating with Node {current_id}: {e}")

def main():
    parser = argparse.ArgumentParser(description="EZ Motion Motor ID and Voltage Configurator")
    parser.add_argument("--current_id", type=int, default=2, help="The current CAN ID of the motor (default: 2)")
    parser.add_argument("--new_id", type=int, default=9, required=True, help="The new CAN ID to assign to the motor")
    parser.add_argument("--channel", type=str, default="can0", help="CAN interface (default: can0)")
    
    args = parser.parse_args()

    configurator = EZMotionConfigurator(channel=args.channel)
    
    try:
        configurator.connect()
        configurator.update_node_id(args.current_id, args.new_id)
    finally:
        configurator.disconnect()

if __name__ == "__main__":
    main()
