#!/usr/bin/env python3
# TODO record and save logs of the data
import bleak
import asyncio
import threading

"""
A class for controlling a collection of PieZense pneumatic systems.
"""
class PieZense:
    millibar=1.0
    Pa=100.0
    kPa=0.1
    bar=0.001
    PSI=0.0145038
    mmHg=0.750062
    mmH2O=10.19716
    _NUS_RX_UUID = "6E400002-B5A3-F393-E0A9-E50E24DCCA9E"  # Nordic UART RX UUID
    def __init__(self, scale_factor: float = 1.0):
        """
        Common scale_factor values:
        - 1.0       : millibar
        - 100.0     : Pa
        - 0.1       : kPa
        - 0.001     : bar (atmosphere)
        - 0.0145038 : PSI
        - 0.750062  : mmHg
        - 10.19716  : mmH2O
        @param scale_factor: scaling factor for pressure readings (default 1.0)
        """
        self._scale_factor = scale_factor
        self._systems = []
        self._reconnect_task = None
        self._update_lock=threading.Lock()
        self._pressures = []
        self._variables = []
        self._callback_function = None
    class _Forward_List:
        """
        A class representing a list of forwarding configurations for a channel, used internally by the PieZense library
        Each channel of a system has a _Forward_List object to store which channel or channels it is controlling.
        """
        def __init__(self):
            self.target_system_num = []
            self.target_channel_num = []
            self.forwarding_function = []
            self.length = 0
        def addForwarding(self, target_system_num: int, target_channel_num: int, forwarding_function):
            self.target_system_num.append(target_system_num)
            self.target_channel_num.append(target_channel_num)
            self.forwarding_function.append(forwarding_function)
            self.length = self.length + 1
        def getForwarding(self, index: int):
            return (self.target_system_num[index], self.target_channel_num[index], self.forwarding_function[index])
        def clearAllForwarding(self):
            self.target_system_num = []
            self.target_channel_num = []
            self.forwarding_function = []
            self.length = 0
        def removeForwarding(self, target_system_num: int, target_channel_num: int):
            for i in range(self.length):
                if self.target_system_num[i] == target_system_num and self.target_channel_num[i] == target_channel_num:
                    # the ith forwarding matches the requested one to remove
                    del self.target_system_num[i]
                    del self.target_channel_num[i]
                    del self.forwarding_function[i]
                    self.length = self.length - 1
                    return
    class _System:
        """
        A class representing a single PieZense system, used internally by the PieZense library
        """
        def __init__(self, system_name: str, channel_count: int):
            self.system_name = system_name
            self.channel_count = channel_count
            self.client = None
            self.connected = False
            self.forwards = [PieZense._Forward_List() for _ in range(channel_count)]
            self.initialized = False
    def addSystem(self, system_name: str, channel_count: int, variable_count: int=1) -> int:
        """ 
        register a PieZense system that you want to connect to

        @param system_name: Bluetooth name of the PieZense system
        @param channel_count: Number of channels in the system, in the future this second argument may become optional
        @return: index of the registered system
        """
        return self._addSystem(system_name, channel_count, variable_count)
    def _addSystem(self, system_name, channel_count, variable_count) -> int:
        self._systems.append(self._System(system_name, channel_count))
        self._pressures.append([0]*channel_count)
        if variable_count > 1:
            self._variables.append(([0]*variable_count)*channel_count)
        return len(self._systems) - 1
    def connect(self):
        """
        start the process of connecting to all registered systems then run forever while connected

        this function is not blocking, it starts a thread in the background

        use isEverythingConnected() to tell when all registered systems have become connected
        """
        self._reconnect_task = threading.Thread(target=lambda: asyncio.run(self._connect()),daemon=True) # True daemon means it runs in background
        self._reconnect_task.start()
 
    def _notification_handler(self, system_index, sender, data):
        expected_without_ts = self._systems[system_index].channel_count * 2
        pressure_data = data[:expected_without_ts]
        self._update_lock.acquire()
        self._systems[system_index].connected = True
        for follower_id in range(self._systems[system_index].channel_count):
            low_byte = pressure_data[follower_id * 2]
            high_byte = pressure_data[follower_id * 2 + 1]
            pressure_value = (high_byte << 8) | low_byte

            if(pressure_value != 65500):
                self._systems[system_index].initialized = True
            else: # it is 65500
                self._systems[system_index].initialized = False
            
            # print(f"[notification_handler] Device={system_index}, Follower={follower_id}, "f"Pressure={pressure_value}")
            
            self._pressures[system_index][follower_id] = pressure_value * self._scale_factor
            # TODO self._variables[system_index][follower_id] = 0  # Placeholder for additional variable

            forward_list=self._systems[system_index].forwards[follower_id]

            for forward_i in range(forward_list.length): # for each follower this channel is controlling
                (target_system_num, target_channel_num, forwarding_function)=forward_list.getForwarding(forward_i)
                pressure_value_scaled=forwarding_function(pressure_value * self._scale_factor)
                self.sendSetpoint(target_system_num, target_channel_num, pressure_value_scaled)

        self._update_lock.release()
                 
        if self._callback_function:
            self._callback_function(system_index, self._pressures[system_index])

    async def _connect(self):
        try:
            while(True):
                for system_i, system in enumerate(self._systems):
                    if not system.client or not system.client.is_connected: # never connected, or disconnected
                        # print(f"Need device: {system.system_name}, connected={system.client.is_connected if system.client else 'N/A'}")
                        if system.client: # client exists but is disconnected, clean up
                            await system.client.disconnect()
                            system.client = None
                            system.connected = False
                            system.initialized = False
                        device = await bleak.BleakScanner.find_device_by_name(system.system_name,timeout=2.0)
                        # print(f"Scanned for device: {system.system_name}")
                        if device: # found device, connect to it
                            # print(f"Found device: {device}")
                            system.client = bleak.BleakClient(device)
                            await system.client.connect()
                            if system.client.is_connected:
                                # print(f"Connected to device: {system.system_name}")
                                services=system.client.services
                                # print(f"Services: {services}")
                                for service in services:
                                    for char in service.characteristics:
                                        if "notify" in char.properties:
                                            # print(f"[{system.system_name}] Subscribing to {char.uuid}")
                                            await system.client.start_notify(
                                                char.uuid,
                                                lambda sender, data, idx=system_i: self._notification_handler(idx, sender, data)
                                            )
                            else:
                                # print(f"Failed to connect to device: {system.system_name}")
                                system.client = None # clean up failed client, will retry next loop
                                system.connected = False
                                system.initialized = False
                    await asyncio.sleep(1)

        except Exception as e:
            print(f"[_connect] Exception in connection loop: {e}")
        finally:
            print("[_connect] Event loop exiting.")
            # for all clients, disconnect
            for system in self._systems:
                try:
                    if system.client and system.client.is_connected:
                        await system.client.disconnect()
                        system.client = None
                except Exception as e:
                    print(f"Error disconnecting from system {system.system_name}: {e}")

    def disconnect(self):
        """
        disconnect from all connected systems, call before and only before exiting your program
        """
        for system in self._systems:
            if system.client and system.client.is_connected:
                asyncio.run(system.client.disconnect())
                system.client = None

    def isEverythingConnected(self) -> bool:
        """
        check if all registered systems are currently connected
        @return: bool: True if all systems are connected, False otherwise
        """
        return all( (system.client and system.client.is_connected and system.connected and system.initialized) for system in self._systems)
    
    async def _sendData(self, system_num: int, data_bytes):
        await self._systems[system_num].client.write_gatt_char(self._NUS_RX_UUID, data_bytes, response=False)

    def sendSetpoint(self, system_num: int, channel_num: int, setpoint: float) -> bool:
        """
        send a pressure setpoint to a channel of a system

        @param system_num: index of the system to send the setpoint to (later this may support system names too)
        @param channel_num: index of the channel to send the setpoint to
        @param setpoint: pressure setpoint
        """
        if self._systems[system_num].initialized == False:
            return False
        if system_num < 0 or system_num >= len(self._systems):
            return False
        if channel_num < 0 or channel_num >= self._systems[system_num].channel_count:
            return False
        pressure_value= int(setpoint / self._scale_factor)
        if self._systems[system_num].client and self._systems[system_num].client.is_connected:
            try:
                data_bytes = bytes([
                    channel_num,
                    (pressure_value >> 8) & 0xFF,
                    pressure_value & 0xFF
                ])
                asyncio.create_task(self._sendData(system_num, data_bytes))

                return True
            except Exception as e:
                print(f"[sendSetpoint] Error sending to system {system_num}, channel {channel_num}: {e}")
                return False
        return False

    def sendSetpointBatch(self, setpoint_batch: list):
        """
        send a batch of pressure setpoints to multiple systems and channels

        @param setpoint_batch: a list of tuples, each containing (system_num (int), channel_num (int), setpoint (float))

        example: [ (0, 0, 1013), (0, 1, 500), (1, 0, 750) ]
        """
        for (system_num, channel_num, setpoint) in setpoint_batch:
            self.sendSetpoint(system_num, channel_num, setpoint)

    def getPressureReadings(self) -> list:
        """
        get the latest pressure readings from all connected systems
        @return: list: a list of lists, where each inner list contains the pressure readings for a system
        """
        with self._update_lock:
            return self._pressures
    
    def getVariableReadings(self,whichVariable:int) -> list:
        """
        get the latest readings of an additional variable from all connected systems
        @param whichVariable: index of the variable to read (0-indexed)
        @return: list: a list of lists, where each inner list contains the whichVariable-numbered variable for each channel of a system
        
        for example, if each system is sending leak-corrected pressure as a second variable, then getVariableReadings(1) will return a list that is like getPressureReadings() but with leak-corrected pressures instead of raw pressures
        """
        pass #TODO

    def setCallback(self, callback_function):
        """
        set a callback function to be called when new pressure data is received
        @param callback_function: a function that takes three arguments: system_num (int) and pressure_data (list)
        """
        self._callback_function = callback_function
    
    def addForwarding(self, source_system_num: int, source_channel_num: int, target_system_num: int, target_channel_num: int, forwarding_function):
        """
        configure pressure forwarding from one channel to another

        @param source_system_num: index of the source system
        @param source_channel_num: index of the source channel
        @param target_system_num: index of the target system
        @param target_channel_num: index of the target channel
        @param forwarding_function: a function that takes a pressure value and returns a modified pressure value (for example lambda x: 4*(x-1100)+1100)
        """
        self._systems[source_system_num].forwards[source_channel_num].addForwarding(target_system_num, target_channel_num, forwarding_function)

    def addForwardingBatch(self, forwarding_batch: list):
        """
        configure multiple pressure forwardings in a batch

        @param forwarding_batch: a list of tuples, each containing (source_system_num (int), source_channel_num (int), target_system_num (int), target_channel_num (int), forwarding_function (function))
        """
        for (source_system_num, source_channel_num, target_system_num, target_channel_num, forwarding_function) in forwarding_batch:
            self.addForwarding(source_system_num, source_channel_num, target_system_num, target_channel_num, forwarding_function)

    def stopForwarding(self, source_system_num: int, source_channel_num: int, target_system_num: int, target_channel_num: int):
        """
        stop pressure forwarding from a channel to a channel

        @param source_system_num: index of the source system
        @param source_channel_num: index of the source channel
        @param target_system_num: index of the target system
        @param target_channel_num: index of the target channel
        """
        self._systems[source_system_num].forwards[source_channel_num].removeForwarding(target_system_num, target_channel_num)

    def clearAllForwarding(self):
        """
        clear all pressure forwarding configurations
        """
        for system in self._systems: # for each system
            for channel_forwards in system.forwards: # for each channel's forwarding list
                channel_forwards.clearAllForwarding()

    # TODO: include list of configuration options with library? maybe a way to get the list and get descriptions of it? Should documentation of config options be part of the PieZense library or should it associated with the follower firmware?
    def sendConfig(self, system_num: int, channel_num: int, config_data: dict):
        """
        configure a channel of a system
        @param system_num: index of the system to configure (later this may support system names too)
        @param channel_num: index of the channel to configure
        @param config_data: configuration data to send

        config_data example {"set_act_mode": 1, "set_pid_Pvalues_p": 0.5}
        """
        if system_num < 0 or system_num >= len(self._systems):
            return False
        if self._systems[system_num].initialized == False:
            return False
        client=self._systems[system_num].client
        if not client or not client.is_connected:
            return False
        if channel_num < 0 or channel_num >= self._systems[system_num].channel_count:
            return False
        try:
            for key, value in config_data.items():
                # key is config parameter name, value is config parameter value
                if not isinstance(key, str):
                    continue
                if key.strip() == "":
                    continue
                if not isinstance(value, (bool, int, float)):
                    continue
                message=f"{key}:{value}"
                messageInBytes = bytes([channel_num]) + message.encode('utf-8')
                asyncio.create_task(self._sendData(system_num, messageInBytes))

            return True
        except Exception as e:
            print(f"[sendConfig] Error sending config to system {system_num}, channel {channel_num}: {e}")
            return False

    def sendConfigToMany(self, which: list, config_data: dict):
        """
        send the same configuration data to multiple systems and channels

        @param which: a list of tuples, each containing (system_num (int), channel_num (int))
        @param config_data: configuration data to send

        example: which = [ (0, 0), (0, 1), (1, 0) ], config_data = {"set_act_mode": 1}
        """
        for (system_num, channel_num) in which:
            self.sendConfig(system_num, channel_num, config_data)

    def sendConfigBatch(self, config_data_batch: list):
        """
        send a batch of configuration data to multiple systems and channels

        @param config_data_batch: a list of tuples, each containing (system_num (int), channel_num (int), config_data (dict))

        example: [ (0, 0, {"set_act_mode": 1}), (0, 1, {"set_act_mode": 0}), (1, 0, {"set_act_mode": 1}) ]
        """
        for (system_num, channel_num, config_data) in config_data_batch:
            self.sendConfig(system_num, channel_num, config_data)

    async def setMode(self, mode: dict):
        """
        Set the operating mode of the PieZense systems.

        @param mode: dictionary describing mode actions

        The mode dict should contain the following keys: 
          - "reset_config": config_data_batch: a list of tuples, each containing (system_num (int), channel_num (int), config_data (dict))
          - "setpoints": setpoint_batch: a list of tuples, each containing (system_num (int), channel_num (int), setpoint (float))
          - "wait_time": seconds, float
          - "forwarding": forwarding_batch: a list of tuples, each containing (source_system_num (int), source_channel_num (int), target_system_num (int), target_channel_num (int), forwarding_function (function))
          - "final_config": config_data_batch: a list of tuples, each containing (system_num (int), channel_num (int), config_data (dict))

        setMode performs the following steps:
          1. Clears all forwarding.
          2. Sends the "reset_config" set of configuration changes (e.g. set actuator mode / default loop params).
          3. Sends the "setpoints" set of pressure setpoints.
          4. Waits for "wait_time" seconds.
          5. Adds the new set of channel forwardings.
          6. Sends the "final_config" set of configuration changes.
        """
        self.clearAllForwarding()
        if "reset_config" in mode:
            self.sendConfigBatch(mode["reset_config"])
        if "setpoints" in mode:
            self.sendSetpointBatch(mode["setpoints"])
        if "wait_time" in mode:
            time_to_wait = mode["wait_time"]
            if isinstance(time_to_wait, (int, float)) and time_to_wait > 0:
                await asyncio.sleep(time_to_wait) # TODO: make this non-blocking?
        if "final_config" in mode:
            self.sendConfigBatch(mode["final_config"])
        if "forwarding" in mode:
            self.addForwardingBatch(mode["forwarding"])

    def makePIDConfig(self, which: list, pP: float, pI: float, pD: float, vP:float,vI:float,vD:float) -> list:
        """
        create a batch of PID configuration data for multiple systems and channels

        @param which: a list of tuples, each containing (system_num (int), channel_num (int))
        @param pP: proportional gain for pressure loop
        @param pI: integral gain for pressure loop
        @param pD: derivative gain for pressure loop
        @param vP: proportional gain for vent loop
        @param vI: integral gain for vent loop
        @param vD: derivative gain for vent loop
        @return: a list of tuples, each containing (system_num (int), channel_num (int), config_data (dict))
        """
        config_batch = []
        for (system_num, channel_num) in which:
            config_data = {
                "set_pid_Pvalues_p": pP,
                "set_pid_Pvalues_i": pI,
                "set_pid_Pvalues_d": pD,
                "set_pid_Vvalues_p": vP,
                "set_pid_Vvalues_i": vI,
                "set_pid_Vvalues_d": vD
            }
            config_batch.append( (system_num, channel_num, config_data) )
        return config_batch
    
    