"""
WHAD Enhanced ShockBurst base connector

This module provides the `ESB` connector class to interact with Enhanced
ShockBurst compatible devices.
"""
from scapy.packet import Packet

from whad.device import Connector
from whad.helpers import message_filter
from whad.exceptions import UnsupportedDomain, UnsupportedCapability

from whad.scapy.layers.esb import ESB_Hdr,ESB_Payload_Hdr,ESB_Ack_Response

from whad.hub.discovery import Domain, Capability
from whad.hub.esb import ESBMetadata, EsbNodeAddress, Commands
from whad.hub.generic.cmdresult import Success, CommandResult
from whad.hub.events import JammedEvt

from whad.esb.esbaddr import ESBAddress

class ESB(Connector):
    """
    Enhanced ShockBurst protocol connector.

    This connector drives a Enhanced ShockBurst (ESB) capable device with
    ESB-specific WHAD messages. It is required by various role classes to
    interact with a real device and pre-process domain-specific messages.
    """

    domain = "esb"

    def format(self, packet):
        """
        Converts a scapy packet with its metadata to a tuple containing a scapy packet with
        the appropriate header and the timestamp in microseconds.
        """
        if isinstance(packet, bytes):
            packet = ESB_Hdr(packet)
        return self.hub.esb.format(packet)

    def __init__(self, device=None, synchronous=False):
        """
        Initialize the connector, open the device (if not already opened), discover
        the services (if not already discovered).

        If `synchronous` is set to True, PDUs must be processed manually and
        won't be forwarded to PDU-related callbacks.
        """
        self.__ready = False
        super().__init__(device)

        # Metadata cache
        self.__cached_channel = None
        self.__cached_address = None

        # Capability cache
        self.__can_send = None
        self.__can_send_raw = None

        # Open device and make sure it is compatible
        self.device.open()
        self.device.discover()

        # Check if device supports Enhanced ShockBurst
        if not self.device.has_domain(Domain.Esb):
            raise UnsupportedDomain("ESB")

        # Mark connector as ready
        self.__ready = True

        # Set synchronous mode
        self.enable_synchronous(synchronous)

    def close(self):
        """Stop the WHAD adpater and close its related device.
        """
        self.stop()
        self.device.close()

    def can_sniff(self):
        """
        Determine if the device implements a sniffer mode.
        """
        commands = self.device.get_domain_commands(Domain.Esb)
        return (
            (commands & (1 << Commands.Sniff)) > 0 and
            (commands & (1 << Commands.Start))>0 and
            (commands & (1 << Commands.Stop))>0
        )

    def can_send(self):
        """
        Determine if the device can transmit packets.
        """
        if self.__can_send is None:
            commands = self.device.get_domain_commands(Domain.Esb)
            self.__can_send = ((commands & (1 << Commands.Send))>0 or
                               (commands & (1 << Commands.SendRaw)))
        return self.__can_send

    def send(self,pdu, address=None, channel=None, retransmission_count=1):
        """
        Send Enhanced ShockBurst packets (on a single channel).
        """
        if self.can_send():
            # If we don't have address or channels, use the cached ones
            tx_address = address if address is not None else self.__cached_address
            tx_channel = channel if channel is not None else self.__cached_channel
            if self.support_raw_pdu():
                # If we support raw PDU but only got a payload, build a packet
                if ESB_Hdr not in pdu:
                    packet = ESB_Hdr(address=tx_address) / pdu
                else:
                    packet = pdu
            # if we don't support raw PDU and got a packet, crop to keep only the payload
            elif ESB_Hdr in pdu:
                packet = pdu[ESB_Payload_Hdr:]
            # if we don't support raw PDU and got a payload, keep it as it is
            else:
                packet = pdu

            # Generate TX metadata
            packet.metadata = ESBMetadata()
            packet.metadata.channel = tx_channel
            packet.metadata.address = tx_address
            packet.metadata.retransmission_count = retransmission_count
            if self.support_raw_pdu():
                packet.metadata.raw = True
                if bytes.fromhex(packet.address[:2])[0] >= 0x80:
                    packet.preamble = 0xAA
                else:
                    packet.preamble = 0x55

            # Send packet
            return super().send_packet(packet)

        # Cannot send, return False.
        return False

    def support_raw_pdu(self):
        """
        Determine if the device supports raw PDU.
        """
        if self.__can_send_raw is None:
            capabilities = self.device.get_domain_capability(Domain.Esb)
            self.__can_send_raw = not capabilities & Capability.NoRawData
        return self.__can_send_raw


    def can_set_node_address(self):
        """
        Determine if the device can configure a Node address.
        """
        commands = self.device.get_domain_commands(Domain.Esb)
        return (
            (commands & (1 << Commands.SetNodeAddress)) > 0
        )

    def start_sniff(self, channel : int = None, address : str = "FF:FF:FF:FF:FF",
              show_acknowledgements : bool = False):
        """
        Sniff Enhanced ShockBurst packets.

        :param channel: Channel to listen, None to iterate over all possible channels
        :type channel: int
        :param address: Device address to target
        :type address: str
        :param show_acknowledgements: Sniff packets acknowledgements if set to True (default: False)
        :type show_acknowledgements: bool
        """
        if not self.can_sniff():
            raise UnsupportedCapability("Sniff")

        self.__cached_address = ESBAddress(address)
        if channel is None:
            # Enable scanning mode
            channel = 0xFF
        else:
            self.__cached_channel = channel

        msg = self.hub.esb.create_sniff_mode(
            EsbNodeAddress(self.__cached_address.value),
            channel,
            show_acknowledgements
        )
        resp = self.send_command(msg, message_filter(CommandResult))
        return isinstance(resp, Success)



    def can_be_prx(self):
        """
        Determine if the device implements a Primary Receiver (PRX) role mode.
        """
        commands = self.device.get_domain_commands(Domain.Esb)
        return (
            (commands & (1 << Commands.PrimaryReceiverMode)) > 0 and
            (commands & (1 << Commands.Start))>0 and
            (commands & (1 << Commands.Stop))>0
        )

    def enable_prx_mode(self, channel):
        """
        Enable Enhanced ShockBurst primary receiver (PRX) mode.
        """
        if not self.can_be_prx():
            raise UnsupportedCapability("PrimaryReceiverMode")

        # Check that we got a channel provided
        if channel is None:
            return False

        # Keep provided channel in cache
        self.__cached_channel = channel

        # Create a PrxMode message.
        msg = self.hub.esb.create_prx_mode(channel)

        resp = self.send_command(msg, message_filter(CommandResult))
        return isinstance(resp, Success)


    def can_be_ptx(self):
        """
        Determine if the device implements a Primary Transmitter (PTX) role mode.
        """
        commands = self.device.get_domain_commands(Domain.Esb)
        return (
            (commands & (1 << Commands.PrimaryTransmitterMode)) > 0 and
            (commands & (1 << Commands.Start))>0 and
            (commands & (1 << Commands.Stop))>0
        )

    def enable_ptx_mode(self, channel):
        """
        Enable Enhanced ShockBurst primary transmitter (PTX) mode.
        """
        if not self.can_be_ptx():
            raise UnsupportedCapability("PrimaryTransmitterMode")


        if channel is None:
            # Enable scanning mode
            channel = 0xFF
        else:
            # Keep channel in cache
            self.__cached_channel = channel

        # Create a PtxMode message.
        msg = self.hub.esb.create_ptx_mode(channel)

        resp = self.send_command(msg, message_filter(CommandResult))
        return isinstance(resp, Success)

    def set_node_address(self, address):
        """
        Enable Enhanced ShockBurst primary receiver (PRX) mode.
        """
        if not self.can_set_node_address():
            raise UnsupportedCapability("SetNodeAddress")

        node_address = ESBAddress(address)
        # Keep address in cache
        self.__cached_address = node_address

        # Create a SetNodeAddress message
        msg = self.hub.esb.create_set_node_address(
            EsbNodeAddress(node_address.value)
        )

        resp = self.send_command(msg, message_filter(CommandResult))
        return isinstance(resp, Success)

    def start(self):
        """
        Start currently enabled mode.
        """
        # Create a Start message.
        msg = self.hub.esb.create_start()

        resp = self.send_command(msg, message_filter(CommandResult))
        return isinstance(resp, Success)

    def stop(self):
        """
        Stop currently enabled mode.
        """
        # Create a Stop message.
        msg = self.hub.esb.create_stop()

        resp = self.send_command(msg, message_filter(CommandResult))
        return isinstance(resp, Success)

    def on_discovery_msg(self, message):
        pass

    def on_generic_msg(self, message):
        pass

    def on_domain_msg(self, domain, message):
        pass

    def on_packet(self, packet):
        """Incoming packet callback.
        """
        if not self.__ready:
            return

        # Dispatch packet
        if packet.metadata.raw:
            self.on_raw_pdu(packet)
        else:
            self.on_pdu(packet)

    def on_event(self, event):
        """Process incoming events.
        """
        if not self.__ready:
            return

        if isinstance(event, JammedEvt):
            self.on_jammed(event.timestamp)

    def on_raw_pdu(self, packet: Packet):
        """Callback to handle incoming raw PDU

        :param packet: Incoming PDU
        :type packet: Packet
        """
        # Extract the PDU from raw packet
        if ESB_Payload_Hdr in packet:
            pdu = packet[ESB_Payload_Hdr:]
        else:
            pdu = ESB_Payload_Hdr()/ESB_Ack_Response()

        # Propagate metadata to PDU
        pdu.metadata = packet.metadata
        self.on_pdu(pdu)

    def on_pdu(self, packet):
        """Callback for incoming PDU.

        For now, all compatible devices report raw PDUs when it comes to ESB,
        so this should not be called.
        """

    def on_jammed(self, timestamp: int):
        """Jammed event handler.
        """
