# Copyright 2025 XMOS LIMITED.
# This Software is subject to the terms of the XMOS Public Licence: Version 1.
from enum import Enum, auto
import sys
import platform
from pathlib import Path
from xscope_endpoint import Endpoint, QueueConsumer
import time

class XscopeControl():
    """
    Class containing functions for sending control commands to the device over xscope.

    An instance of this class is created when entering the runtime context of XcoreAppControl, which
    also starts the DUT xrun process, and the functions in this class can be called as follows:

    with XcoreAppControl(adapter_id, xe_name, verbose=verbose) as xcoreapp:
        xcoreapp.xscope_host.xscope_controller_cmd_connect()
        xcoreapp.xscope_host.xscope_controller_cmd_set_dut_macaddr(0, "01:02:03:04:05:06")

    It can be used standalone without having to create an instance of XcoreAppControl first.
    When using this class in a standalone manner, do the following:
    In one terminal, run the dut app using xrun --xscope-port. For example:

    xrun --xscope-port localhost:12340 <xe file>

    Then in another terminal, from the lib_ethernet/tests directory, open a python shell and run

    from xscope_host import XscopeControl
    xscope_host = XscopeControl("localhost", "12340", verbose=True)

    Follow this with the commands you wish to send to the device. For example,

    xscope_host.xscope_controller_cmd_connect()
    xscope_host.xscope_controller_cmd_set_host_macaddr("a4:ae:12:77:86:97")
    xscope_host.xscope_controller_cmd_set_dut_macaddr(0, "01:02:03:04:05:06")
    xscope_host.xscope_controller_cmd_set_dut_tx_packets(0, 10000, 345)
    xscope_host.xscope_controller_cmd_shutdown()

    """
    class XscopeCommands(Enum):
        """
        Class containing supported commands that can be sent to the device over xscope.
        Extend the list below when adding a new command. The include file for the cmds enum included when compiling the device test apps
        is autogenerated from the list below in the CMakeLists for the test apps. The XscopeCommands.write_to_h_file() method is called for this.
        """
        CMD_DEVICE_SHUTDOWN = auto()
        CMD_SET_DEVICE_MACADDR = auto()
        CMD_SET_HOST_MACADDR = auto()
        CMD_HOST_SET_DUT_TX_PACKETS = auto()
        CMD_SET_DUT_RECEIVE = auto()
        CMD_DEVICE_CONNECT = auto()
        CMD_EXIT_DEVICE_MAC = auto()
        CMD_SET_DUT_TX_SWEEP = auto()

        """
        Method for generating a .h file containing the enum defined above, which can be included
        when compiling the device test apps that access this enum when responding to control
        commands over xscope.
        """
        @classmethod
        def write_to_h_file(cls, filename):
            filename = Path(filename)
            dir_path = filename.parent
            dir_path.mkdir(parents=True, exist_ok=True)

            with open(filename, "w") as fp:
                name = filename.name
                name = name.replace(".", "_")
                fp.write(f"#ifndef __{name}__\n")
                fp.write(f"#define __{name}__\n\n")
                fp.write("typedef enum {\n")
                for member in cls:
                    fp.write(f"\t{member.name} = {member.value},\n")
                fp.write("}xscope_cmds_t;\n\n")
                fp.write("#endif\n")

    def __init__(self, host, port, timeout=30, verbose=False):
        """
        Initialise the XscopeControl() class.

        This saves the arguments passed to __init__ into class variables.
        Note that this does not create the xscope endpoint and connects to the device.
        """
        self.host = host
        self.port = port
        self.timeout = timeout
        self.verbose = verbose
        self._ep = None

    def xscope_controller_do_command(self, cmds, connect=True):
        """
        Runs the xscope host app to connect to the device and execute a command over xscope port

        Parameters:
        cmds (list): byte list containing the command + arguments for the command that needs to be executed.
        The command is executed by,
        connect to the xscope server to communicate with the device
        send cmd + args bytes to the device over xscope.
        receive ack from the device over an xscope probe
        disconnect from the xscope server.

        connect (bool, optional, default=True): Whether or not to connect to the xscope server when running a command.
        For most commands, connecting, communicating and disconnecting works and is the default behaviour.
        However, for some special cases, such as reading data being sent from the device over a probe while sending commands,
        it is required to connect once, and send multiple commands.

        For eg.
        # connect to the xscope server and start reading timestamps recorded on a probe
        xscope_host.xscope_controller_start_timestamp_recorder()

        # send cmd to device to transmit num_packets of a given packet len. Note that the host app is already connected
        # to the xscope server at this point, since the probe capture needs to start before the device starts transmitting packets
        # in order to record timestamps from the first packet itself.
        xcoreapp.xscope_host.xscope_controller_cmd_set_dut_tx_packets(0, num_packets, packet_len, connect=False)

        time.sleep(5) # wait while device transmits.

        # Disconnect from the xscope server and return the probed timestamps
        probe_timestamps = xcoreapp.xscope_host.xscope_controller_stop_timestamp_recorder()

        Returns:
        device stdout from executing the command
        """
        if connect:
            ep = Endpoint()
        else:
            ep = self._ep
        probe = QueueConsumer(ep, "command_ack")

        if connect:
            if ep.connect(hostname=self.host, port=self.port):
                print("Xscope Host app failed to connect")
                assert False
        if self.verbose:
            print(f"Sending {cmds} bytes to the device over xscope")
        ep.publish(bytes(cmds))
        ack = probe.next()
        if self.verbose:
            print(f"Received ack {ack}")

        device_stdout = ep._captured_output.getvalue() # stdout from the device
        if self.verbose:
            print("stdout from the device:")
            print(device_stdout)

        if connect:
            ep.disconnect()
        if ack == None:
            print("Xscope host received no response from device")
            print(f"device stdout: {device_stdout}")
            assert False
        return device_stdout


    def xscope_controller_cmd_connect(self):
        """
        Run command to ensure that the device is set up and ready to communicate via ethernet

        Returns:
        device stdout from executing the command
        """
        ret = self.xscope_controller_do_command([XscopeControl.XscopeCommands['CMD_DEVICE_CONNECT'].value])
        # This is to wait for debugger's link up in case its in the path. Ideally we should be doing
        # dbg.wait_for_links_up(): but for non debugger tests there's no guarantee that the debugger is in the path
        # TODO Why is this still needed even after waiting for debugger links to be up??
        time.sleep(2)
        return ret


    def xscope_controller_cmd_shutdown(self):
        """
        Run command to shutdown the device application threads and exit

        Returns:
        device stdout from executing the command
        """
        return self.xscope_controller_do_command([XscopeControl.XscopeCommands['CMD_DEVICE_SHUTDOWN'].value])

    def xscope_controller_cmd_set_dut_macaddr(self, client_index, mac_addr):
        """
        Run command to set the source mac address of a client running on the device.

        Parameters:
        client_index (int8): index of the client.
        mac_addr (str, eg. "11:e0:24:df:33:66"): mac address

        Returns:
        device stdout from executing the command
        """
        mac_addr_bytes = [int(i, 16) for i in mac_addr.split(":")]
        cmd_plus_args = [XscopeControl.XscopeCommands['CMD_SET_DEVICE_MACADDR'].value, client_index]
        cmd_plus_args.extend(mac_addr_bytes)
        return self.xscope_controller_do_command(cmd_plus_args)

    def xscope_controller_cmd_set_host_macaddr(self, mac_addr):
        """
        Run command to inform the device of the host's mac address. This is required for a TX client running on the device to know the destination
        mac address for the ethernet packets it is transmitting.

        Parameters:
        mac_addr (str, eg. "f0:f1:f2:f3:f4:f5"): mac address

        Returns:
        device stdout from executing the command
        """
        mac_addr_bytes = [int(i, 16) for i in mac_addr.split(":")]
        cmd_plus_args = [XscopeControl.XscopeCommands['CMD_SET_HOST_MACADDR'].value]
        cmd_plus_args.extend(mac_addr_bytes)
        return self.xscope_controller_do_command(cmd_plus_args)

    def xscope_controller_cmd_set_dut_tx_packets(self, client_index, arg1, arg2, connect=True):
        """
        Run command to inform the TX clients on the device, information about the packets it needs to transmit.

        The paramters are different depending on whether the TX client in high or low priority

        Parameters:
        client_index (int32): index of the client
        arg1 (int32): number of packets to send for LP client, QAV BW in bps for HP client
        arg2 (int32): packet payload length in bytes
        connect (bool, optional, default=True): When running this command, whether or not to initiate a new connection to
        the xscope server

        Returns:
        device stdout from executing the command
        """
        cmd_plus_args = [XscopeControl.XscopeCommands['CMD_HOST_SET_DUT_TX_PACKETS'].value]
        for a in [client_index, arg1, arg2]: # client_index, arg1 and arg2 are int32
            bytes_to_append = [(a >> (8 * i)) & 0xFF for i in range(4)]
            cmd_plus_args.extend(bytes_to_append)
        return self.xscope_controller_do_command(cmd_plus_args, connect=connect)


    def xscope_controller_cmd_set_dut_receive(self, client_index, recv_flag):
        """
        Run command to get an RX client on the device to start or stop receiving packets.

        Parameters:
        client_index (int8): RX client index on the device
        recv_flag (bool): Flag indicating whether to receive (1) or not receive (0) packets

        Returns:
        device stdout from executing the command
        """
        cmd_plus_args = [XscopeControl.XscopeCommands['CMD_SET_DUT_RECEIVE'].value, client_index, recv_flag]
        return self.xscope_controller_do_command(cmd_plus_args)

    def xscope_controller_cmd_restart_dut_mac(self):
        """
        Run command to restart the device ethernet Mac.

        Returns:
        device stdout from executing the command
        """
        return self.xscope_controller_do_command([XscopeControl.XscopeCommands['CMD_EXIT_DEVICE_MAC'].value])

    def xscope_controller_cmd_set_dut_tx_sweep(self, client_index, connect=True):
        """
        Run command to get a TX client on the device to transmit packets sweeping through all valid (64 - 1518 bytes) packet lengths,
        sending a fixed (hardcoded in the application) number of packets for each packet length.

        Parameters:
        client_index (int8): RX client index on the device
        connect (bool, optional, default=True): When running this command, whether or not to initiate a new connection to
        the xscope server

        Returns:
        device stdout from executing the command
        """
        cmd_plus_args = [XscopeControl.XscopeCommands['CMD_SET_DUT_TX_SWEEP'].value, client_index]
        return self.xscope_controller_do_command(cmd_plus_args, connect=connect)

    def xscope_controller_start_timestamp_recorder(self):
        """
        Connect to the xscope server and start recording TX timestamps sent by the device over the 'tx_start_timestamp' probe.
        """
        ep = Endpoint()
        probe = QueueConsumer(ep, "tx_start_timestamp")

        if ep.connect(hostname=self.host, port=self.port):
            print("Xscope Host app failed to connect")
            assert False

        self._ep = ep
        self._probe = probe


    def xscope_controller_stop_timestamp_recorder(self):
        """
        Disconnect from device and return all data collected from self._probe

        Returns:
        probe_output (list) : data captured over the probe in the form of an array of bytes
        """
        print(f"{self._probe.queue.qsize()} elements in the queue")
        probe_output = []
        for i in range(self._probe.queue.qsize()):
            probe_output.extend(self._probe.next())
        device_stdout = self._ep._captured_output.getvalue() # stdout from the device
        print("stdout from the device:")
        print(device_stdout)
        self._ep.disconnect()
        return probe_output



"""
WARNING:
Do not change the main function since it's called from CMakeLists.txt to autogenerate the xscope commands enum .h file
"""
if __name__ == "__main__":
    print("Generate xscope cmds enum .h file")
    assert len(sys.argv) == 2, ("Error: filename not provided" +
                    "\nUsage: python generate_xscope_cmds_enum_h_file.py <.h file name, eg. enum.h>\n")

    XscopeControl.XscopeCommands.write_to_h_file(sys.argv[1])


