import socket
import struct
import subprocess
import numpy as np
import serial

probe = {
        '10mV': np.int8(0),
        '20mV': np.int8(1), 
        '50mV': np.int8(2), 
        '100mV': np.int8(3), 
        '200mV': np.int8(4), 
        '500mV': np.int8(5), 
        '1V': np.int8(6), 
        '2V': np.int8(7), 
        '5V': np.int8(8),
        '10V': np.int8(9), 
        '20V': np.int8(10), 
        '50V': np.int8(11), 
        '100V': np.int8(12), 
        '200V': np.int8(13)
        }

class adc_default:
    def __init__(self, port):
        self.port = port
        self.ndata = []
        self.channels_data = []
        self.measure_code = (0x00, 0)

        self.proto = {
            'open': 0x01,
            'set_rate': 0x07,
            'set_points': 0x06,
            'config_channels': 0x09,
            'set_premeasure': 0x28,
            'set_trignum': 0x19,
            'start': 0x3B,
            'stop': 0x03,
        }

    def connect(self, port=0):
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.client_socket.settimeout(5.0)
        resp = self.client_socket.connect(('127.0.0.1', self.port))
        return resp

    def resp_handler(self, resp, cmd):
        if(resp[0] != 0xAA):
            return (-1, 0)
        if(resp[1] == 0xFF):
            return (0xFF, resp[3])
        return (resp[1], 0)
            
    def open(self):
        msg = struct.pack('<BB', 0xAA, self.proto['open'])
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        return self.resp_handler(resp, self.proto['open'])
    
    def set_rate(self, rate):
        print('send sample_rate')
        msg = struct.pack('<BBI', 0xAA, self.proto['set_rate'], np.uint32(rate))
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        return resp

    def set_points(self, points):
        print('send points')
        msg = struct.pack('<BBI', 0xAA, self.proto['set_points'], len(points))
        for p in points:
            msg += p.to_bytes(length=4, byteorder='little')
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        return self.resp_handler(resp, self.proto['set_points'])

    def config_channels(self, nchannels, channel_ranges, trig_channel, trig_dirrection, trig_threshold, auto_trigger_time):
        print('configure')
        print(channel_ranges)
        channel_ranges_bytes = b''
        for r in channel_ranges:
            channel_ranges_bytes += r.to_bytes(length=1)
        print(channel_ranges_bytes)

        msg = struct.pack(f'<BBI{nchannels}sBiHh', 0xAA, self.proto['config_channels'], nchannels, channel_ranges_bytes, trig_channel, trig_dirrection, trig_threshold, auto_trigger_time)
        print(msg[6:8])
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        return self.resp_handler(resp, self.proto['config_channels'])

    def set_premeasure(self, trig_premeasure):
        print('set pre-measure')
        msg = struct.pack('<BBI', 0xAA, self.proto['set_premeasure'], trig_premeasure)
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        return self.resp_handler(resp, self.proto['set_premeasure'])

    def set_trignum(self, trig_num):
        print('set trignum')
        msg = struct.pack('<BBI', 0xAA, self.proto['set_trignum'], trig_num)
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        return self.resp_handler(resp, self.proto['set_trignum'])
    
    def start(self):
        data = b''
        temp = b''
        print('start')
        msg = struct.pack('<BB', 0xAA, self.proto['start'])
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        while((resp[1] != 0xFB) and (resp[1] != 0xFF)):
            print(f"Data confirmed with code {resp[1]}")
            msg = struct.pack('<BB', 0xAA, 0x3A)
            self.client_socket.send(msg)
            resp = self.client_socket.recv(4096)
        msg = struct.pack('<BB', 0xAA, 0x3B)
        self.client_socket.send(msg)

        resp = self.client_socket.recv(4096)
        nsignal = 0
        nchannel = 0
        last_signal = 0
        last_channel = 0
        while(resp[1] != 0xFC):
            if(resp[1] == 0xFF):
                self.measure_code = self.resp_handler(resp, self.proto['start'])
                print('error')
                return (self.resp_handler(resp, self.proto['start']), self.ndata)
            elif(resp[1] == 0xFD):
                magic, cmd, nsignal, nchannel, total_packets, npacket, ndata_sent, first = struct.unpack(f'<BBIIIIIh', resp[0:24])
                eod = ndata_sent+22
                print(f'sig {nsignal}, ch {nchannel}, packet {npacket} / {total_packets}')
                temp = resp[22:eod]
                #print(f'First: {first} or First: {struct.unpack('<h', temp[0:2])}')
                msg = struct.pack('<BB', 0xAA, 0x3D)
                self.client_socket.send(msg)

                if(nchannel > last_channel):
                    print('append channel')
                    self.channels_data.append(data)
                    data = b''
                    last_channel += 1
                if(nsignal > last_signal):
                    print('append channel')
                    self.channels_data.append(data)
                    data = b''
                    last_channel = 0
                    last_signal += 1
                    print('append data')
                    self.ndata.append(self.channels_data)
                    self.channels_data = []
                data += temp
            resp = self.client_socket.recv(4096)
            print(f'RESP: {resp[1]}')
        msg = struct.pack('<BB', 0xAA, 0x3C)
        self.client_socket.send(msg)
        print('append channel')
        self.channels_data.append(data)
        print('append data')
        self.ndata.append(self.channels_data)
        self.measure_code = self.resp_handler(resp, self.proto['start'])
        return (self.resp_handler(resp, self.proto['start']), self.ndata)
    
    def stop(self):
        print('set post-measure')
        msg = struct.pack('<BB', 0xAA, self.proto['stop'])
        self.client_socket.send(msg)
        #resp = self.client_socket.recv(4096)
        return self.resp_handler((0x00, 0x00), self.proto['stop'])
    
    def clear_ndata(self):
        self.ndata = []
        self.channels_data = []

    
class sync_default:
    def __init__(self, port):
        if(port != -1):
            self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        self.proto = {
            'upload': 0x01,
            'trig_wait': 0x02
        }
        
        self.serial = None

    def connect(self, port):
        self.client_socket.settimeout(5.0)
        resp = self.client_socket.connect(('localhost', port))
        return resp

    def resp_handler(self, resp, cmd):
        if(resp[0] != 0xAA):
            return (-1, 0)
        if(resp[1] == 0xFF):
            return (0xFF, resp[3])
        return (resp[1], 0)
            
    def upload(self, programPath, filePath, port):
        try:
            process = subprocess.run([programPath, filePath, '-p', str(port), '--debug'])
        except subprocess.CalledProcessError as e:
            print(f"Ошибка команды {e.cmd}!")
        return process.returncode
    
    def serial_open(self, port):
        try:
            self.serial = serial.Serial('COM' + str(port), 9600, timeout=1)
        except serial.SerialException as e:
            print(e)
            return (0xFF, -1)
        return(0x00, 0)

    def trig_wait(self, port):
        try:
            self.serial.write(b'e')
        except serial.SerialException as e:
            print(e)
            return (0xFF, -1)
        return(0x00, 0)

    def serial_close(self):
        try:
            self.serial.close()
        except:
            return (0xFF, -1)
        return (0x00, 0)

class sdr_default:
    def __init__(self, port):
        if(port != -1):
            self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        self.proto = {
            'transf': 0x01
        }

    def connect(self, port):
        self.client_socket.settimeout(5.0)
        resp = self.client_socket.connect(('localhost', port))
        return resp

    def transf(self, programPath, filePath, freq, rate, ampl, gain):
        try:
            process = subprocess.run([programPath, '-t', filePath, '-f', str(freq), '-s', str(rate), '-a', str(ampl), '-x', str(gain)], check=True)
        except subprocess.CalledProcessError as e:
            print(f"Ошибка команды {e.cmd}!")
            return (0xFF, -1)
        return (0x00, process.returncode)
    
class gra_default:
    def __init__(self, ip, port):
        if(port != -1):
            self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.client_socket.settimeout(5.0)
            self.client_socket.connect((ip, port))

        self.proto = {
            'reset': 0x0001,
            'ps_on': 0x0003,
            'ps_off': 0x0004,
            'state': 0x0005,
            'upload': 0x0006,
            'send_packet': 0x0007,
            'confirm': 0x0008
        }

    def connect(self, port):
        self.client_socket.settimeout(5.0)
        resp = self.client_socket.connect(('localhost', port))
        return resp

    def reset(self):
        msg = struct.pack('<HH', 0xAAAA, self.proto['reset'])
        self.client_socket.send(msg)

    def ps_on(self):
        msg = struct.pack('<HH', 0xAAAA, self.proto['ps_on'])
        self.client_socket.send(msg)

    def ps_off(self):
        msg = struct.pack('<HH', 0xAAAA, self.proto['ps_off'])
        self.client_socket.send(msg)

    def ps_off(self):
        msg = struct.pack('<HH', 0xAAAA, self.proto['ps_off'])
        self.client_socket.send(msg)

    def state(self):
        msg = struct.pack('<HH', 0xAAAA, self.proto['state'])
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        if(len(resp) == 6):
            magic, cmd, state = struct.unpack('<HHH', resp)
            return (cmd, state, 0, 0)
        elif(len(resp) > 6):
            magic, cmd, state, errbits, errpoint, current_setting_amp, current_amp, first_sensor, second_sensor = struct.unpack('<HHHIIiiii', resp)
            return (cmd, state, errbits, errpoint)
        
    def send_packet(self, nom, nodes_num, nodes_buf):
        i = 0
        time1, time2 = struct.unpack('<HH', nodes_buf[i][0].to_bytes(4, 'little'))
        nom = nom | 0x8000
        msg = struct.pack('<HHHHHHH', 0xAAAA, self.proto['send_packet'], nom, nodes_num, time1, time2, nodes_buf[i][1].to_bytes(2, 'little'))
        while(i < nodes_num):
            i += 1
            time1, time2 = struct.unpack('<HH', nodes_buf[i][0].to_bytes(4, 'little'))
            msg += struct.pack('<HHH', time1, time2, nodes_buf[i][1].to_bytes(2, 'little'))
        self.client_socket.send(msg)
        resp = self.client_socket.recv(4096)
        data_resp = struct.unpack('<HHH', resp)
        return data_resp[3]
    
    def upload(self, points, nodes):
        points1, points2 = struct.unpack('<HH', points.to_bytes(4, 'little'))
        nodes1, nodes2 = struct.unpack('<HH', nodes[points-1][0].to_bytes(4, 'little'))
        msg = struct.pack('<HHHHHH', 0xAAAA, self.proto['upload'], points1, points2, nodes1, nodes2)
        self.client_socket.send(msg)
        seg_num = 0
        total_packets = points // 200 + 1
        first_idx = 0
        while(seg_num < total_packets):
            if((points - first_idx) < 200):
                nodes_packet = points - first_idx
            else:
                nodes_packet = 200
            last_idx = first_idx + nodes_packet
            resp = self.send_packet(seg_num, nodes_packet, nodes[first_idx:last_idx])
            print(resp)
            first_idx = last_idx
            seg_num += 1
        msg = struct.pack('<HH', 0xAAAA, 0x0008)
        self.client_socket.send(msg)