import struct
import threading
import time
from modbus_tk.modbus_tcp import TcpServer
import modbus_tk.defines as cst

class wtMBtcp:
    def __init__(self, ip='127.0.0.1', port=502, salve_id = 1):
        self.ip = ip
        self.port = port
        self.slave_id = salve_id
        self.slave = []
        self.server = None
        self.network_status = False
        threading.Thread(target=self.initialize_data_store).start()
    

    def all_keys_exist(self, variable, required_keys):
        return all(key in variable for key in required_keys)
    
    def initialize_data_store(self):
        try:
            if self.network_status is False:
                self.network_status = True
                self.server = TcpServer(address=self.ip, port=self.port)
                for index, id in enumerate(self.slave_id):
                    self.slave.append(self.server.add_slave(id))  
                    self.slave[index].add_block("A", cst.HOLDING_REGISTERS, 0, 10)  # 模拟100个寄存器
                self.server.start()
        except Exception as exp:
            self.network_status = False
            for index, id in enumerate(self.slave_id):
                self.slave[index].remove_block("A")
            self.server.stop()
            self.slave = []
            print(exp) 
        time.sleep(10)
                         
    def modify_register(self, id, cmd, address, new_value, data_type):
        try:
            for index, id_cn in enumerate(self.slave_id):
                if id_cn == id:
                    if data_type == 'int': 
                        # int 类型使用两个寄存器，大端模式  # AB  
                        values = struct.unpack('>H', struct.pack('>H', new_value))
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'int1':     
                        # int 类型使用两个寄存器，小端模式  # BA
                        values = struct.unpack('>H', struct.pack('<H', new_value))
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'float':
                        # float 类型使用四个寄存器，大端模式 # AB CD
                        values = struct.unpack('>HH', struct.pack('>f', new_value))
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'float2':
                        # float 类型使用四个寄存器，小端模式 # CD BA
                        values = struct.unpack('<HH', struct.pack('<f', new_value))
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'float3':
                        values = struct.unpack('<HH', struct.pack('>f', new_value)) # BA DC
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'float4':
                        values = struct.unpack('>HH', struct.pack('<f', new_value)) # DC AB
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'long':
                        # long 类型使用四个寄存器，大端模式  # AB CD
                        values = struct.unpack('>HH', struct.pack('>i', new_value))
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'long2':
                        # long 类型使用四个寄存器，小端模式  # CD BA
                        values = struct.unpack('<HH', struct.pack('<i', new_value))
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'long3':
                        values = struct.unpack('<HH', struct.pack('>i', new_value))  # BA DC
                        self.slave[index].set_values("A", address, list(values))
                    elif data_type == 'long4':
                        values = struct.unpack('>HH', struct.pack('<i', new_value))  # DC AB
                        self.slave[index].set_values("A", address, list(values))
        except Exception as exp:
            print(exp)    

    def reconnect_tcp(self):
        self.server.stop()
        self.server = TcpServer(address=self.ip, port=self.port)
        self.server.start()
        self.slave = self.server.add_slave(self.slave_id) 
