import pyvisa as visa
import serial
import socket


class Connection:
    
    def write(self, cmd):
        pass

    def query(self, cmd):
        pass

    def read_raw(self):
        pass

    def write_raw(self, data):
        pass

    def set_visa_attribute(self, attr, value):
        pass


class VisaConnection(Connection):

    def __init__(self):
        self.device = None

    def open(self, rsrc):
        rm = "" # r"C:\windows\system32\rsvisa32.dll"
        self.device = visa.ResourceManager(rm).open_resource(rsrc)
        self.device.timeout = 2000
        self.resource = rsrc
        if "SOCKET" in rsrc.upper():
            self.device.read_termination = '\n'


    def write(self, cmd):
        self.device.write(cmd)

    def query(self, cmd):
        return self.device.query(cmd)

    def read_raw(self, chunk=100):
#        return self.device.read_raw(chunk) // does not work
        mes = self.device.visalib.read(self.device.session, chunk)
        return mes[0]


    def write_raw(self, data):
        self.device.write_raw(data)

    def close(self):
        self.device.close()

    def set_visa_attribute(self,attr, value):
        self.device.set_visa_attribute(attr, value)




class CdcConnection(Connection):

    def __init__(self):
        self.device = None


    def open(self, Port):
        self.device = serial.Serial()
        self.device.baudrate = 115200 * 4
        self.device.port = f"{Port}"
        self.device.timeout = 0.1
        self.device.open()
        assert(self.device.is_open)

    def write(self, cmd):
        self.device.write(bytes(f"{cmd}", "utf-8"))

    def write_raw(self, data):
        self.device.write(data)

    def query(self, cmd):
        self.device.write(bytes(f"{cmd}\n", "utf-8"))
        rc = self.device.readline(100)
        s = rc.decode("utf-8")
        return s

    def read_raw(self, chunk=100):
        rc = self.device.readline(chunk)
        return rc
    
    def read(self ):
        return self.device.read()
    
    def close(self):
        self.device.close()


class SocketConnection(Connection):

    def __init__(self):
        self.device = None


    def open(self, Host, Port):
        self.device = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 
        self.device.connect((Host, Port))


    def write(self, cmd):
        self.device.sendall(bytes(f"{cmd}\n", "utf-8"))

    def write_raw(self, data):
        self.device.sendall(data)

    def query(self, cmd):
        self.device.sendall(bytes(f"{cmd}\n", "utf-8"))
        data = self.device.recv(1024)
        return data.decode('utf-8')

    def read_raw(self, chunk=100):
        rc = self.device.recv(chunk)
        return rc
    
    def close(self):
        self.device.close()
