import serial
import time
import struct
import random
import threading

class SimpleSTS3215:
    """
    A lightweight, thread-safe driver for STS3215 servos on clean Half-Duplex UART.
    
    This driver handles the communication protocol for Feetech STS3215 serial bus servos.
    It is designed for "clean" hardware setups (e.g., dedicated half-duplex adapter or
    properly configured Pi UART) where the transmitted data is NOT echoed back to the RX pin.
    
    Attributes:
        ser (serial.Serial): The underlying pyserial object.
        lock (threading.RLock): Reentrant lock to ensure atomic transactions.
    """
    
    # Memory Addresses
    ADDR_ID = 5
    ADDR_BAUD_RATE = 6
    ADDR_MIN_ANGLE_LIMIT = 9
    ADDR_MAX_ANGLE_LIMIT = 11
    ADDR_MODE = 33
    ADDR_TORQUE_ENABLE = 40
    ADDR_ACC = 41
    ADDR_GOAL_POSITION = 42
    ADDR_GOAL_TIME = 44
    ADDR_GOAL_SPEED = 46
    ADDR_PRESENT_POSITION = 56
    ADDR_PRESENT_SPEED = 58
    ADDR_PRESENT_LOAD = 60
    ADDR_PRESENT_VOLTAGE = 62
    ADDR_PRESENT_TEMP = 63
    ADDR_MOVING = 66
    
    def __init__(self, port='/dev/ttyAMA0', baudrate=1000000, timeout=0.5):
        """
        Initializes the serial connection to the servo bus.

        Args:
            port (str): The serial port path (e.g., '/dev/serial0' or 'COM3').
            baudrate (int): Communication speed in bps. Default is 1,000,000.
            timeout (float): Read timeout in seconds. Default 0.5s is conservative for Pi Zero.
        """
        self.ser = serial.Serial(
            port=port,
            baudrate=baudrate,
            timeout=timeout,
            bytesize=serial.EIGHTBITS,
            stopbits=serial.STOPBITS_ONE,
            parity=serial.PARITY_NONE,
            rtscts=False,
            dsrdtr=False
        )
        self.header = [0xFF, 0xFF]
        self.lock = threading.RLock() # Reentrant lock for nested calls

    def _calc_checksum(self, data):
        """
        Calculates the standard STS protocol checksum.
        
        Formula: ~(Sum of all parameters) & 0xFF

        Args:
            data (list[int]): List of bytes to checksum.

        Returns:
            int: The calculated checksum byte.
        """
        total = sum(data)
        return (~total) & 0xFF

    def _write_packet(self, servo_id, instruction, params):
        """
        Builds and sends a communication packet to the servo.
        
        Includes retry logic for WRITE instructions that expect an ACK.
        
        Args:
            servo_id (int): The target servo ID (0-253) or Broadcast (254).
            instruction (int): The protocol instruction ID.
            params (list[int]): A list of data bytes to send as parameters.
        """
        # Determine if we expect a response (Write instruction to unicast ID)
        expect_ack = (servo_id != 0xFE) and (instruction == 3)
        retries = 3 if expect_ack else 1

        for attempt in range(retries):
            with self.lock:
                # Length = Instruction(1) + Params(N) + Checksum(1)
                length = len(params) + 2
                content = [servo_id, length, instruction] + params
                checksum = self._calc_checksum(content)
                packet = self.header + content + [checksum]
                
                # Clear buffer before sending to ensure we read OUR response
                self.ser.reset_input_buffer()
                self.ser.write(bytearray(packet))
                
                if expect_ack:
                    try:
                        self._read_packet(0)
                        return # Success
                    except (TimeoutError, ValueError):
                        if attempt == retries - 1:
                            raise # Propagate error on last attempt
                        # If failed, loop will retry the write
                        time.sleep(0.01) # Brief pause before retry
                else:
                    # For non-ack packets (like read requests handled by caller), return immediately
                    return

    def _read_packet(self, expected_params_len):
        """
        Reads and validates a response packet from the serial bus.
        
        Includes robust header searching to handle occasional line noise.

        Args:
            expected_params_len (int): The number of data bytes expected in the payload.

        Returns:
            bytes: The payload parameters extracted from the packet.
        """
        # Header(2) + ID(1) + Len(1) + Err(1) + Params(N) + Checksum(1)
        total_len = 6 + expected_params_len
        
        # Read enough bytes to cover the packet. 
        data = self.ser.read(total_len)
        
        if len(data) < total_len:
            raise TimeoutError(f"Expected {total_len} bytes, got {len(data)}: {data.hex().upper()}")

        # Safety: Check for Header. If data[0] isn't 0xFF, we might have noise.
        # We try to slide the window if we have extra bytes, or just fail if strict.
        while len(data) >= total_len:
            if data[0] == 0xFF and data[1] == 0xFF:
                break
            else:
                # Shift buffer by 1 byte and try to read one more
                next_byte = self.ser.read(1)
                if not next_byte:
                    raise ValueError(f"Invalid Header: {data.hex().upper()}")
                data = data[1:] + next_byte
        
        if data[0] != 0xFF or data[1] != 0xFF:
             raise ValueError(f"Invalid Header after search: {data.hex().upper()}")
             
        error = data[4]
        if error != 0:
            error_desc = []
            if error & 1: error_desc.append("Voltage")
            if error & 2: error_desc.append("Angle")
            if error & 4: error_desc.append("Overheat")
            if error & 8: error_desc.append("Range")
            if error & 16: error_desc.append("Checksum")
            if error & 32: error_desc.append("Overload")
            print(f"[WARN] Servo Error: {error} ({', '.join(error_desc)})")

        return data[5:-1] # Return just the parameters

    def ping(self, servo_id):
        """Checks if a servo is present."""
        for _ in range(3):
            try:
                # Manual write packet construction to handle the Read logic purely here
                length = 2
                instruction = 1 # Ping
                content = [servo_id, length, instruction]
                checksum = self._calc_checksum(content)
                packet = self.header + content + [checksum]
                
                with self.lock:
                    self.ser.reset_input_buffer()
                    self.ser.write(bytearray(packet))
                    self._read_packet(0)
                return True
            except (TimeoutError, ValueError):
                time.sleep(0.01)
        return False

    def scan(self, maxid=253):
        """Scans IDs 0-maxid and returns a list of found servos."""
        found = []
        print("Scanning bus for servos...")
        # Reduce timeout temporarily for faster scan
        old_timeout = self.ser.timeout
        self.ser.timeout = 0.05
        
        for i in range(maxid):
            if self.ping(i):
                print(f"Found Servo ID: {i}")
                found.append(i)
        
        self.ser.timeout = old_timeout
        return found

    def change_id(self, current_id, new_id):
        """
        Changes the ID of a connected servo.
        
        WARNING: If multiple servos have 'current_id', they will ALL change.
        Connect one servo at a time to use this safely.
        """
        print(f"Changing ID {current_id} to {new_id}...")
        # Unlock EEPROM
        self._write_packet(current_id, 3, [55, 0]) 
        # Write ID (Addr 5)
        self._write_packet(current_id, 3, [self.ADDR_ID, new_id])
        # Lock EEPROM
        self._write_packet(new_id, 3, [55, 1])
        print("Done.")

    def set_position(self, servo_id, position, speed=2400, acc=50):
        """Moves the servo to a specific target position."""
        position = int(position)
        position = max(0, min(3*4095, position))
        
        pos_L, pos_H = position & 0xFF, (position >> 8) & 0xFF
        spd_L, spd_H = speed & 0xFF, (speed >> 8) & 0xFF
        
        # We assume _write_packet handles retries for individual packets.
        # However, for set_position, we want the group to succeed.
        # Since _write_packet retries internally, this simple sequence is robust.
        with self.lock:
            self._write_packet(servo_id, 3, [self.ADDR_ACC, acc])
            self._write_packet(servo_id, 3, [self.ADDR_GOAL_SPEED, spd_L, spd_H])
            self._write_packet(servo_id, 3, [self.ADDR_GOAL_POSITION, pos_L, pos_H])

    def set_speed(self, servo_id, speed):
        """
        Sets the target speed for the servo.

        In Position Mode: Limits the maximum speed for the trajectory.
        In Speed Mode (Wheel Mode): Sets the continuous rotation speed and direction.

        Args:
            servo_id (int): The ID of the servo.
            speed (int): Target speed in steps/sec. 
                         - Positive: Clockwise (usually).
                         - Negative: Counter-Clockwise (Bit 15 set).
                         - Range: ~ -3400 to +3400.
        """
        # STS Protocol uses Bit 15 for direction, Bits 0-14 for magnitude
        direction_bit = 0
        if speed < 0:
            speed = abs(speed)
            direction_bit = (1 << 15)
            
        speed = min(speed, 4095) # Clamp magnitude if needed
        
        final_speed_val = speed | direction_bit
        
        spd_L = final_speed_val & 0xFF
        spd_H = (final_speed_val >> 8) & 0xFF
        
        with self.lock:
             self._write_packet(servo_id, 3, [self.ADDR_GOAL_SPEED, spd_L, spd_H])

    def get_position(self, servo_id):
        """Reads the current absolute position of the servo shaft."""
        for _ in range(3):
            try:
                # Instruction 2 (Read)
                length = 4 
                instruction = 2
                params = [self.ADDR_PRESENT_POSITION, 2]
                content = [servo_id, length, instruction] + params
                checksum = self._calc_checksum(content)
                packet = self.header + content + [checksum]
                
                with self.lock:
                    self.ser.reset_input_buffer()
                    self.ser.write(bytearray(packet))
                    data = self._read_packet(2)
                return (data[1] << 8) | data[0]
            except (TimeoutError, ValueError):
                time.sleep(0.01)
        raise TimeoutError(f"Failed to read position from ID {servo_id}")

    def get_feedback(self, servo_id):
        """Retrieves a comprehensive status snapshot of the servo."""
        for _ in range(3):
            try:
                length = 4
                instruction = 2
                params = [self.ADDR_PRESENT_POSITION, 11]
                content = [servo_id, length, instruction] + params
                checksum = self._calc_checksum(content)
                packet = self.header + content + [checksum]
                
                with self.lock:
                    self.ser.reset_input_buffer()
                    self.ser.write(bytearray(packet))
                    d = self._read_packet(11)
                
                def to_signed(low, high):
                    val = (high << 8) | low
                    if val > 32767: val -= 65536
                    return val

                return {
                    "position": (d[1] << 8) | d[0],
                    "speed": to_signed(d[2], d[3]),
                    "load": to_signed(d[4], d[5]),
                    "voltage": d[6] / 10.0,
                    "temp": d[7],
                    "moving": d[10]
                }
            except (TimeoutError, ValueError):
                time.sleep(0.01)
        raise TimeoutError(f"Failed to read feedback from ID {servo_id}")
        
    def torque_enable(self, servo_id, enable=True):
        """Enables or disables the motor torque."""
        val = 1 if enable else 0
        self._write_packet(servo_id, 3, [self.ADDR_TORQUE_ENABLE, val])

    def set_mode(self, servo_id, mode):
        """Sets the operational mode of the servo."""
        self.torque_enable(servo_id, False)
        self._write_packet(servo_id, 3, [self.ADDR_MODE, mode])

    def configure_multiturn(self, servo_id):
        """
        Configures the servo for multi-turn mode (unlimited rotation).
        
        Sets both the Min Angle Limit (Addr 9) and Max Angle Limit (Addr 11) to 0.
        In the STS3215 protocol, setting both limits to 0 disables the single-turn 
        angle safety check.
        
        Args:
            servo_id (int): The ID of the servo.
        """
        print(f"Configuring ID {servo_id} for multi-turn mode...")
        # Unlock EEPROM
        self._write_packet(servo_id, 3, [55, 0])
        # Write 0 to Min Angle Limit (Addr 9)
        self._write_packet(servo_id, 3, [self.ADDR_MIN_ANGLE_LIMIT, 0, 0])
        # Write 0 to Max Angle Limit (Addr 11)
        self._write_packet(servo_id, 3, [self.ADDR_MAX_ANGLE_LIMIT, 0, 0])
        # Lock EEPROM
        self._write_packet(servo_id, 3, [55, 1])
        print(f"ID {servo_id} configured. Full signed 16-bit range enabled.")

    def reset_encoder_to_midpoint(self, servo_id):
        """
        Resets the current hardware position to 2048 (Midpoint).
        
        Uses the special Torque Enable command (128) to set current position as offset.
        Useful when installing motors at unknown angles. 
        Call this before calibration routine to avoid hitting 0 limit.
        """
        print(f"Resetting ID {servo_id} current position to 2048...")
        # Unlock EEPROM (required for some models, but 128 is usually a RAM op)
        self._write_packet(servo_id, 3, [55, 0]) 
        # Write 128 (0x80) to Torque Enable
        self._write_packet(servo_id, 3, [self.ADDR_TORQUE_ENABLE, 128])
        # Lock EEPROM
        self._write_packet(servo_id, 3, [55, 1])
        # Re-enable torque after reset if needed (usually defaults to off/safe)
        self.torque_enable(servo_id, True)

if __name__ == "__main__":
    sts = SimpleSTS3215(port='/dev/serial0', timeout=0.5) 

    # --- CLI Mode: Set ID ---
    if len(sys.argv) == 3 and sys.argv[1] == "setid":
        try:
            target_new_id = int(sys.argv[2])
            if not (0 <= target_new_id <= 253):
                raise ValueError("ID must be 0-253")
        except ValueError as e:
            print(f"Invalid ID: {e}")
            sys.exit(1)

        print(f"Attempting to set single connected servo to ID {target_new_id}...")
        
        # Perform a full scan (0-253) to ensure we don't accidentally broadcast to multiple motors
        # If we see ANY more than 1 motor, we abort to prevent ID collisions.
        servos = sts.scan() 
        
        if len(servos) == 0:
            print("Error: No servos found. Check power and connections.")
        elif len(servos) > 1:
            print(f"Error: Found {len(servos)} servos ({servos}).")
            print("SAFETY ABORT: Changing ID with multiple servos connected will change ALL of them to the same ID.")
            print("Please connect ONLY the specific servo you want to configure.")
        else:
            current_id = servos[0]
            print(f"Found exactly one servo at ID {current_id}.")
            
            if current_id == target_new_id:
                print(f"Servo is already set to ID {target_new_id}. No action needed.")
            else:
                sts.change_id(current_id, target_new_id)

    # --- Default Mode: Scan and Test IDs 1-8 ---
    else:
        servos = sts.scan(8)
        
        if not servos:
            print("No servos found. Check power/wiring.")
            print("Usage: python simple_sts3215.py")
            print("       python simple_sts3215.py setid <new_id>")
            exit()
            
        print(f"Active Servos: {servos}")

        for target_id in servos:
            print(f"Testing Servo ID {target_id}...")
            sts.torque_enable(target_id, True)
            
            pos = sts.get_position(target_id)
            print(f"Current Pos: {pos}")
            
            target_pos = random.randint(500, 3500)
            print(f"Moving to {target_pos}...")
            sts.set_position(target_id, target_pos)
            
            for _ in range(20): 
                try:
                    feedback = sts.get_feedback(target_id)
                    print(f"Pos: {feedback['position']} | Load: {feedback['load']} | Moving: {feedback['moving']}")
                    if abs(feedback['position'] - target_pos) < 15:
                        print("Target Reached!")
                        break
                except Exception as e:
                    print(f"Feedback read error: {e}")
                time.sleep(0.1)
            
            sts.torque_enable(target_id, False)