import bluetooth
import struct
import time

# BLE UUID definitions (Nordic UART Service)
_UART_UUID = bluetooth.UUID("6E400001-B5A3-F393-E0A9-E50E24DCCA9E")
_UART_TX = (
    bluetooth.UUID("6E400003-B5A3-F393-E0A9-E50E24DCCA9E"),
    bluetooth.FLAG_NOTIFY,
)
_UART_RX = (
    bluetooth.UUID("6E400002-B5A3-F393-E0A9-E50E24DCCA9E"),
    bluetooth.FLAG_WRITE | bluetooth.FLAG_WRITE_NO_RESPONSE,
)
_UART_SERVICE = (_UART_UUID, (_UART_TX, _UART_RX))

# BLE IRQ constants
_IRQ_CENTRAL_CONNECT = 1
_IRQ_CENTRAL_DISCONNECT = 2
_IRQ_GATTS_WRITE = 3

class BLEUART:
    """A helper class to manage BLE UART (Serial Over BLE) for raw text data exchange."""
    
    def __init__(self, ble=None, name="ESP32-S3-RLCD"):
        self._ble = ble if ble is not None else bluetooth.BLE()
        self._ble.active(True)
        self._ble.irq(self._irq)
        
        # Register UART service
        ((self._tx_handle, self._rx_handle),) = self._ble.gatts_register_services((_UART_SERVICE,))
        
        self._connections = set()
        self._rx_callback = None
        self._name = name
        
        self._advertise()

    def _irq(self, event, data):
        # Handle connection, disconnection, and writes
        if event == _IRQ_CENTRAL_CONNECT:
            conn_handle, _, _ = data
            self._connections.add(conn_handle)
            print(f"BLE Central connected (handle: {conn_handle})")
            
        elif event == _IRQ_CENTRAL_DISCONNECT:
            conn_handle, _, _ = data
            if conn_handle in self._connections:
                self._connections.remove(conn_handle)
            print(f"BLE Central disconnected (handle: {conn_handle})")
            # Restart advertising
            self._advertise()
            
        elif event == _IRQ_GATTS_WRITE:
            conn_handle, value_handle = data
            if value_handle == self._rx_handle:
                # Read received data
                data_received = self._ble.gatts_read(self._rx_handle)
                if self._rx_callback:
                    try:
                        decoded = data_received.decode("utf-8").strip()
                        self._rx_callback(decoded)
                    except Exception as e:
                        # Fallback for binary / non-UTF-8 data
                        self._rx_callback(data_received)

    def _advertise(self):
        # Generate advertising payload (flags + service UUID = 3 + 18 = 21 bytes)
        adv_data = self._build_advertising_payload(services=[_UART_UUID], include_flags=True)
        # Generate scan response payload (device name, no flags = 2 + len(name) bytes)
        resp_data = self._build_advertising_payload(name=self._name, include_flags=False)
        
        self._ble.gap_advertise(100000, adv_data=adv_data, resp_data=resp_data)
        print(f"BLE advertising started as '{self._name}'")

    def _build_advertising_payload(self, name=None, services=None, include_flags=True):
        """Generates standard BLE advertising packet format."""
        payload = bytearray()
        
        def append(ad_type, value):
            payload.append(len(value) + 1)
            payload.append(ad_type)
            payload.extend(value)
            
        # Flags: General discoverable mode, BR/EDR not supported
        if include_flags:
            append(0x01, struct.pack("B", 0x02 | 0x04))
        
        # Name
        if name:
            append(0x09, name.encode("utf-8"))
            
        # Services UUID
        if services:
            for uuid in services:
                # UUID bytes are in little-endian order
                append(0x07, bytes(uuid))
                
        return bytes(payload)

    def on_rx(self, callback):
        """Register a callback function to handle incoming BLE messages.
        
        Args:
            callback (function): Function accepting a single string argument.
        """
        self._rx_callback = callback
        return callback

    def write(self, data):
        """Sends data (string or bytes) to all connected BLE Centrals."""
        if not self.is_connected():
            return False
            
        # Ensure it is bytes
        if isinstance(data, str):
            data = data.encode("utf-8")
            
        try:
            # Write value to local GATT DB first
            self._ble.gatts_write(self._tx_handle, data)
            # Notify all connected clients
            for conn_handle in self._connections:
                self._ble.gatts_notify(conn_handle, self._tx_handle, data)
            return True
        except Exception as e:
            print(f"Error sending BLE data: {e}")
            return False

    def is_connected(self):
        """Returns True if at least one central is connected."""
        return len(self._connections) > 0

    def close(self):
        """Disconnect and stop BLE."""
        self._ble.gap_advertise(None) # Stop advertising
        for conn_handle in self._connections:
            self._ble.gap_disconnect(conn_handle)
        self._ble.active(False)
        print("BLE closed.")