
from DevicePackages.defaultConfigForDevices import LOGGER, defaultConfig
import pyvisa
import traceback
import time
from PyQt6 import QtCore

def checkIfBussy(func):
    def wrapper(*args, **kwargs):
        #args[0] is self from Connection Class
        if args[0].busy_flag == False:
            args[0].busy_flag = True
            result = func(*args, **kwargs)
            args[0].busy_flag = False
            return result
        else:
            args[0].log.log(f"{args[0].config_dictionary['DEVICENAME']} - Hardware is bussy, try again later","ERROR")
            return "BUSSY"
    return wrapper


class signalsClass(QtCore.QObject):
    error_signal = QtCore.pyqtSignal(str,object)

class ConnectionClass():
    def __init__(self,log:LOGGER = LOGGER("ALL"),\
        config_dictionary:dict = defaultConfig,*args, **kwargs):#args = multiple input variables, **kvargs multiple keyword variables (dictionaries, touples etc.)
    ################ INIT Function ##################
        super().__init__(*args, **kwargs)
        self.log = log
        self.signals = signalsClass()
        self.config_dictionary = config_dictionary
        self.log.log(f" {self.config_dictionary['DEVICENAME']} - STARTING CONNECTION CLASS","UNNECESSARY")
        self.log.log(f" OBJECT: - {self}","UNNECESSARY")
        self._visa_instrument = None
        self._visa_RM = None
        self.busy_flag = False #to prevent multiple attempts to communicate with hardware in the same time from another thread
        self._last_exception = "UNKNOWN"


    def _communicationNotSupported(self):
        self.log.log(f"{self.config_dictionary['DEVICENAME']} - Communication type not supported","ERROR")
        self._last_exception = "Communication type not supported"
        self.signals.error_signal.emit(self._last_exception,self)

    def _handleResult(self, status,result):
        if status == True or result == "VIRTUAL":
            if result == None or result == "":
                self.log.log(f"Status is True, but result: {result} is empty","WARNING")
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - Operation successful","UNNECESSARY")
        
        else:
            self.log.log(f"SENDING ERROR SIGNAL - Operation failed","ERROR")
            self._last_exception = str(result)
            self.signals.error_signal.emit(self._last_exception,self)

    def _exceptionTraceback(self, e):
        self.log.log(f"{self.config_dictionary['DEVICENAME']} - {traceback.format_exc()}","ERROR",oneline=False)
        return  str(e)

    @checkIfBussy
    def connectionSetup(self):
        if self.config_dictionary["COMMUNICATION"] == "VISA":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - connectionSetup as VISA device","HARDWARE")
            status,result = self._visaSetup()
            self._handleResult(status, result)
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - status:{status}, result: {result}","HARDWARE")
            return status,result

        elif self.config_dictionary["COMMUNICATION"] == "VIRTUAL":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - connectionSetup as virtual device","WARNING")
            return True,"VIRTUAL"

        else:
            self._communicationNotSupported()
            return False, "Not Supported"

    @checkIfBussy
    def connect(self):
        if self.config_dictionary["COMMUNICATION"] == "VISA":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - connect to VISA device","HARDWARE")
            status,result = self._visaConnect()
            self._handleResult(status,result)
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - status:{status}, result: {result}","HARDWARE")
            return status,result

        elif self.config_dictionary["COMMUNICATION"] == "VIRTUAL":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - connection as virtual device","WARNING")
            return True, "VIRTUAL"
        else:
            self._communicationNotSupported()
            return False, "Not Supported"

    @checkIfBussy
    def disconnect(self):
        if self.config_dictionary["COMMUNICATION"] == "VISA":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - disconnect from VISA device","HARDWARE")
            status,result = self._visaDisconnect()
            self._handleResult(status,result)
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - status:{status}, result: {result}","HARDWARE")
            return status,result

        elif self.config_dictionary["COMMUNICATION"] == "VIRTUAL":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - disconnecting virtual device","WARNING")
            return True, "VIRTUAL"

        else:
            self._communicationNotSupported()
            return False, "Not Supported"

    @checkIfBussy
    def sendMessage(self,command:str):
        if self.config_dictionary["COMMUNICATION"] == "VISA":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - sendMessage {command}","HARDWARE")
            status,result = self._visaSendMessage(command)
            self._handleResult(status,result)
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - status:{status}, result: {result}","HARDWARE")
            return status,result

        elif self.config_dictionary["COMMUNICATION"] == "VIRTUAL":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - sending virtual message","WARNING")
            return True, "VIRTUAL"

        else:
            self._communicationNotSupported()
            return False, "Not Supported"

    @checkIfBussy
    def readMessage(self):
        if self.config_dictionary["COMMUNICATION"] == "VISA":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - read Message from VISA device","HARDWARE")
            status,result = self._visaReadMessage()
            self._handleResult(status,result)
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - status:{status}, result: {result}","HARDWARE")
            return status,result

        elif self.config_dictionary["COMMUNICATION"] == "VIRTUAL":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - reading virtual message","WARNING")
            return True, "VIRTUAL"

        else:
            self._communicationNotSupported()
            return False, "Not Supported"

    @checkIfBussy
    def queryMessage(self,command:str):
        if self.config_dictionary["COMMUNICATION"] == "VISA":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - queryMessage: {command}","HARDWARE")
            status,result = self._visaQueryMessage(command)
            self._handleResult(status,result)
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - status:{status}, result: {result}","HARDWARE")
            return status,result

        elif self.config_dictionary["COMMUNICATION"] == "VIRTUAL":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - querying virtual message","WARNING")
            return True, "VIRTUAL"

        else:
            self._communicationNotSupported()
            return False, "Not Supported"

    @checkIfBussy
    def blockingSelfCheck(self):
        if self.config_dictionary["COMMUNICATION"] == "VISA":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - blocking SelfCheck VISA device","HARDWARE")
            status,result = self._visaBlockingSelfCheck()
            self._handleResult(status,result)
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - status:{status}, result: {result}","HARDWARE")
            return status,result

        elif self.config_dictionary["COMMUNICATION"] == "VIRTUAL":
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - blocking self check virtual device","WARNING")
            return True, "VIRTUAL"

        else:
            self._communicationNotSupported()
            return False, "Not Supported"
    
    def _visaBlockingSelfCheck(self):
        """Send self check command and wait for response"""
        try:
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - Blocking self check","UNNECESSARY")
            if self._visa_instrument == None:
                self.log.log(f"{self.config_dictionary['DEVICENAME']} - Visa instrument not connected","ERROR")
                return False, "Visa instrument not connected"

            self._visa_instrument.clear()
            status, message = self._visaQueryMessage(self.config_dictionary["SELF_CHECK_COMMAND"])
            if message in self.config_dictionary['SELF_CHECK_RESPONSE']:
                self.log.log(f"{self.config_dictionary['DEVICENAME']} - Blocking self check passed","UNNECESSARY")
                return True, message

            else:
                self.log.log(f"{self.config_dictionary['DEVICENAME']} - Blocking self check failed with message: {message}","ERROR")
                return False, f"Blocking self check failed with message: {message}"
        
        except Exception as e:
            return False, self._exceptionTraceback(e)

    def _visaSetup(self) -> bool:
        """Look at avaiable VISA resources and compare them with the one in the config file"""
        self.log.log(f" {self.config_dictionary['DEVICENAME']} - Visa SETUP","INFO")
        self._visa_RM = pyvisa.ResourceManager()
        for resources in self._visa_RM.list_resources():
            if self.config_dictionary['VISASTRING'] in resources:
                self.log.log(f"{self.config_dictionary['DEVICENAME']} - Visa resource found","INFO")
                return True, "Visa resource found"

        
        self.log.log(f"{self.config_dictionary['DEVICENAME']} - Visa resource NOT found","WARNING")
        self.log.log(f"{self.config_dictionary['DEVICENAME']} Avaiable resources: {self._visa_RM.list_resources()}","INFO")
        self.log.log(f"{self.config_dictionary['DEVICENAME']} If the resource is not found you can still try to connect!","WARNING")
        return False , "resource is not found you can still try to connect!"
        
    def _visaConnect(self) -> bool:
        """Try to connect to the visa resource with settings from self.config_dictionary"""
        try:
            if self._visa_RM == None:
                self._visaSetup()
            self.log.log(f"{self.config_dictionary['DEVICENAME']} - Visa Connecting to {self.config_dictionary['VISASTRING']}","HARDWARE")
            self._visa_instrument:pyvisa.Resource = self._visa_RM.open_resource(self.config_dictionary['VISASTRING'])
            #there is no clean way to set DTR register (restart at connection) in pyvisa -> consume first message from arrigo which ends with \n
            if "BAUDRATE" not in self.config_dictionary: #Connection is propably TCP/IP
                self._visa_instrument.baud_rate = 115200
                self.log.log(f"{self.config_dictionary['DEVICENAME']} - Baudrate not specified in config file, using default 115200","WARNING")
                self.log.log(f"{self.config_dictionary['DEVICENAME']} - If Visa String is not using baud rate, this setting will be ignored","WARNING")
            else:
                self._visa_instrument.baud_rate = self.config_dictionary['BAUDRATE']
            self._visa_instrument.query_delay = float(self.config_dictionary['QUERYDELAY'])
            self._visa_instrument.read_termination = self.config_dictionary['READTERMINATION']
            self._visa_instrument.write_termination = self.config_dictionary['WRITETERMINATION']
            self._visa_instrument.timeout = float(self.config_dictionary['TIMEOUT'])
            time.sleep(0.1)#give instrument some time
            self._visa_instrument.clear()#clear buffers
            return True , "Visa connected"
                
        except Exception as e:
            return False, self._exceptionTraceback(e)


    def _visaDisconnect(self):
        """Disconnect from the visa resource"""
        self.log.log(f"{self.config_dictionary['DEVICENAME']} - Disconnecting from VISA","HARDWARE")
        try:
            self._visa_instrument.close()
            self._visa_RM.close()
            self._visa_instrument = None
            self._visa_RM = None
            return True , "Visa disconnected"
        except Exception as e:
            return False, self._exceptionTraceback(e)

    def _visaSendMessage(self,command:str):
        """Send a command to the visa resource"""
        try:
            self._visa_instrument.write(command)
            return True , f"Sent to visa: {command}"
        except Exception as e:
            return False, self._exceptionTraceback(e)


    def _visaReadMessage(self):
        """Read a message from the visa resource"""
        try:
            return True, self._visa_instrument.read()
        except Exception as e:
            return False, self._exceptionTraceback(e)

    def _visaQueryMessage(self,command:str):
        """Send a query to the visa resource"""
        try:
            return True,self._visa_instrument.query(command)
        except Exception as e:
            return False, self._exceptionTraceback(e)