import socket
import struct

class ModbusRobot:
    def __init__(self, ip="192.168.18.201", port=502):
        self.ip = ip
        self.port = port

    def send_modbus_tcp(self, register, value):
        transaction_id = 0x01
        protocol_id = 0x00
        length = 0x06
        unit_id = 0x01
        function_code = 0x06
        reg_address = register
        reg_value = value
        
        packet = struct.pack('>HHHBBHH', 
                             transaction_id, 
                             protocol_id, 
                             length, 
                             unit_id, 
                             function_code, 
                             reg_address, 
                             reg_value)
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(2.0)
                s.connect((self.ip, self.port))
                print(f"Sending: {packet.hex(' ')}")
                s.sendall(packet)
                response = s.recv(1024)
                return response
        except Exception as e:
            print(f"Error setting: {e}")

    def send_modbus_tcp_32bit(self, start_register, value):
        transaction_id = 0x01
        protocol_id = 0x00
        unit_id = 0x01
        
        low_word = (value >> 16) & 0xFFFF
        high_word = value & 0xFFFF
        function_code = 0x10
        quantity_of_registers = 2
        byte_count = 4
        length = 11

        packet = struct.pack('>HHHBBHHBHH', 
                             transaction_id, 
                             protocol_id, 
                             length, 
                             unit_id, 
                             function_code, 
                             start_register * 2, 
                             quantity_of_registers, 
                             byte_count,
                             high_word, 
                             low_word)
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(2.0)
                s.connect((self.ip, self.port))
                print(f"Send 32-bit ({value}): {packet.hex(' ')}")
                s.sendall(packet)
                response = s.recv(1024)
                print(f"Answer: {response.hex(' ')}")
                return response
        except Exception as e:
            print(f"Error setting: {e}")

    def read_modbus_tcp_registers(self, register, count):
        transaction_id = 0x02
        protocol_id = 0x00
        length = 0x06
        unit_id = 0x01
        function_code = 0x04
        
        packet = struct.pack('>HHHBBHH', 
                             transaction_id, 
                             protocol_id, 
                             length, 
                             unit_id, 
                             function_code, 
                             register, 
                             count)
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(2.0)
                s.connect((self.ip, self.port))
                
                print(f"send request (FC04): {packet.hex(' ')}")
                s.sendall(packet)
                response = s.recv(1024)
                
                if len(response) < 9:
                    print("Error: short message")
                    return None
                
                byte_count = response[8]
                register_data = response[9:]
                num_registers = byte_count // 2
                values = struct.unpack('>' + 'H' * num_registers, register_data)
                
                print(f"answer success: {response.hex(' ')}")
                print(f"Values кregistr (с {register}): {values}")
                return values
        except Exception as e:
            print(f"Error while reading: {e}")