from HCI_tester import *

import copy, time

def CS_setup_cmds(handle):
    return cmd_seq([
            '01 8D 20 05 ' + hp(s.pack('<H', handle)) + ' 03 01 0A', # LE CS Set Default Settings
         #   '01 8A 20 02 80 00', # LE CS Read Remote Support Capabilities
            ])

def CS_terminate_cmd(addr):
    addr_bytes = hp(s.pack('>Q', addr))[6:]
    return cmd_seq([
        '01 0D 20 19 C8 00 C8 00 00 00 ' + addr_bytes + ' 00 F0 00 F0 00 00 00 E8 03 00 01 FF', # LE Create Connection
        ])

# LE CS Create Config
def CS_create_config(handle, id, main_mode=2, sub_mode=1, rtt_type=0, min_main=2, max_main=6, main_mode_rep=1, 
                     n_mode0=3, use_3C=False, ChMap_repeat=1, hat_shape=True, ch3c_jump=3, sync_phy=1):
    return [hex_read(x) for x in [
        '01 90 20 1C ' 
              + hp(s.pack('<HBBBBBBBB', handle, id, 1, main_mode, sub_mode, min_main, max_main, main_mode_rep, n_mode0))
              + ' 00 ' 
              + hp(s.pack('BB', rtt_type, sync_phy)) + ' FC FF 7F FC FF FF FF FF FF 1F '
              + hp(s.pack('BB', ChMap_repeat, 1 if use_3C else 0)) + ' '
              + hp(s.pack('BB', 0 if hat_shape else 1, ch3c_jump))
              + ' 00' 
        #'01 90 20 1C 80 00 00 01 02 01 02 06 01 03 00 00 01 FC FF 7F FC FF FF FF FF FF 1F 03 01 00 03 00', # LE CS Create Config
        ]]

# LE CS Set Procedure Parameters
def CS_set_proc_params(handle, id, with_SNR, proc_count=1, ACI=1, acl_phy=1, subevent_len=30000, proc_interval=20):
    sub = hp(bits24(subevent_len))
    len_byte = '17 ' if with_SNR else '15 '
    SNR_bytes = ' FF FF' if with_SNR else ''
    return [hex_read(x) for x in [
        '01 93 20 ' + len_byte + hp(s.pack('<HB', handle, id)) + ' D0 02 '
            + hp(s.pack('<HHH', proc_interval, proc_interval, proc_count))
            + ' ' + sub + ' ' + sub + ' ' + hp(s.pack('<BB', ACI, acl_phy)) + ' 00 01'
            + SNR_bytes, 
        ]]

def CS_procedure_enable(handle, id, enable=True):
    return [hex_read(x) for x in [
        '01 94 20 04 ' + hp(s.pack('<HB', handle, id)) +  (' 01' if enable else ' 00'), # LE CS Procedure Enable
        ]]

def add_override_data(I, AAs, channels, use_3C, hat_shape, ch3c_jump, markers, fewer_chans, channel_map=None):
    override = bytearray()
    word = 0
    count = 0
    if channels is not None and len(channels) > 0:
        word |= 1
        count += 1 + len(channels)
        override += s.pack('B', len(channels))
        for c in channels:
            override += s.pack('B', c)
    else:
        count += 13
        if channel_map is not None:
            if len(channel_map) != 10:
                raise ValueError("Invalid channel map, should be 10 bytes. channel_map = {}".format(channel_map))
            for c in channel_map:
                override += s.pack('B', c)
        elif fewer_chans:
            override += s.pack('BBBBBBBBBB', 0x0, 0x0, 0x0, 0x0, 0x0, 0xFF, 0xFF, 0xFF, 0xFF, 0x00)
        else:
            override += s.pack('BBBBBBBBBB', 0xFC, 0xFF, 0x7F, 0xFC, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x1F)

        override += s.pack('B', 1 if use_3C else 0)
        override += s.pack('B', 0 if hat_shape else 1)
        override += s.pack('B', ch3c_jump)

    if AAs is not None and len(AAs) == 2:
        word |= 1<<5
        count += 8
        override += s.pack('<II', AAs[0], AAs[1])

    if markers is not None and len(markers) == 3:
        word |= (1<<6) | (1<<7)
        count += 4
        pos0, pos1, val = markers
        override += s.pack('<BBB', pos0, pos1, val)

    I += s.pack('<HB', word, count)
    I += override

def CS_Test_dec23(main_mode=2, sub_mode=0xFF, main_mode_rep=0, mode0_steps=3, rtt_type=0, phy=1, antenna=1, subevent_len=4000000, subevent_interval=0,
                  max_num_subevents=0, TX_power=0, T_IP1=145, T_IP2=145, T_FCS=150, T_PM=40, T_SW=0, Tone_Antenna_Config=0, Companion_Signal_Enable=0,
                  Nonce=0, Channel_Map_Repetition=1, channels=list(range(2,23)) + list(range(26,77)), AAs=(0xC9D07A33,0x2F6C3998),
                  use_3C=False, hat_shape=True, ch3c_jump=3, fewer_chans=False, markers = None):
    '''
    Dec 2023
    STRU_S(HCI_CS_Test_Dec_2023_t) {
        HCI_Main_Mode_Type_e Main_Mode_Type;
        HCI_Sub_Mode_Type_e Sub_Mode_Type;
        DEC_INT_W_MIN_MAX(U8, Main_Mode_Repetition, 0, 3);
        DEC_INT_W_MIN_MAX(U8, Mode_0_Steps, 1, 3);
        HCI_Role_e Role;
        HCI_RTT_Type_e RTT_Type;
        HCI_CS_SYNC_PHY_e CS_SYNC_PHY;
        HCI_CS_SYNC_Test_Antenna_Selection_e CS_SYNC_Test_Antenna_Selection;
        TIME_FIELD_ARR_W_MIN_MAX(U8, CS_Subevent_Len, 1us, 3, 1250, 4000000);
        TIME_ENUM(HCI_CS_Subevent_Interval_e, CS_Subevent_Interval, 625us);
        ENUM_W_MIN_MAX(HCI_CS_Test_Max_Num_Subevents_E, Max_Num_Subevents, 1, 20);
        ENUM_W_UNITS(HCI_Transmit_Power_Level_e, Transmit_Power_Level, dBm);
        HCI_T_IP1_time_e T_IP1_time;
        HCI_T_IP2_time_e T_IP2_time;
        HCI_T_FCS_time_e T_FCS_time;
        HCI_T_PM_time_e T_PM_time;
        HCI_T_SW_time_e T_SW_time;
        HCI_Tone_Antenna_Config_e Tone_Antenna_Config;
        HCI_Companion_Signal_Enable_Test_e Companion_Signal_Enable_Test;
        U8 DRBG_Nonce[2];
        INT_W_MIN_MAX(U8, Channel_Map_Repetition, 1, 0xFF);
        HCI_Override_Config_t Override_Config;
        U8 PRINT_DEC(Override_Parameters_Length);
        U8 Override_Parameters_Data[1];

        STRUCT_MIN_MAX(28, 65535)
        STRUCT_UI_TREE_OUTPUT_FNS;
        NO_AUTOGEN_UI_FNS;
    }
    STRU_E_VLA1(HCI_CS_Test_Dec_2023_t);
    '''
    I = bytearray()
    I += s.pack('<BHB', 1, 0x2095, 0)
    I += s.pack('<8B', main_mode, sub_mode, main_mode_rep, mode0_steps, 0, rtt_type, phy, antenna)
    I += bits24(subevent_len)
    I += s.pack('<HBb6B', subevent_interval, max_num_subevents, TX_power, T_IP1, T_IP2, T_FCS, T_PM, T_SW, Tone_Antenna_Config)
    I += s.pack('<BHB', Companion_Signal_Enable, Nonce, Channel_Map_Repetition)

    add_override_data(I, AAs, channels, use_3C, hat_shape, ch3c_jump, markers, fewer_chans)

    I[3] = len(I) - 4

    R = copy.deepcopy(I)
    R[8] = 1 # Set Role

    return I, R


def CS_Test_BT6(main_mode=2, sub_mode=0xFF, main_mode_rep=0, mode0_steps=3, rtt_type=0, phy=1, antenna=1, subevent_len=4000000, subevent_interval=0,
                max_num_subevents=0, TX_power=0, T_IP1=145, T_IP2=145, T_FCS=150, T_PM=40, T_SW=0,Tone_Antenna_Config=0, Companion_Signal_Enable=0,
                Nonce=0, Channel_Map_Repetition=1, channels=list(range(2,23)) + list(range(26,77)), AAs=(0xC9D07A33,0x2F6C3998), SNRs = (0xFF, 0xFF),
                use_3C=False, hat_shape=True, ch3c_jump=3, fewer_chans=False, markers = None, channel_map = None):
    '''
    STRU_S(HCI_CS_Test_Atlanta_r04_t) {
        HCI_Main_Mode_Type_e Main_Mode_Type;
        HCI_Sub_Mode_Type_e Sub_Mode_Type;
        DEC_INT_W_MIN_MAX(U8, Main_Mode_Repetition, 0, 3);
        DEC_INT_W_MIN_MAX(U8, Mode_0_Steps, 1, 3);
        HCI_CS_Role_e CS_Role;
        HCI_RTT_Type_e RTT_Type;
        HCI_CS_SYNC_PHY_e CS_SYNC_PHY;
        HCI_CS_SYNC_Test_Antenna_Selection_e CS_SYNC_Test_Antenna_Selection;
        TIME_FIELD_ARR_W_MIN_MAX(U8, CS_Subevent_Len, 1us, 3, 1250, 4000000);
        TIME_ENUM(HCI_CS_Subevent_Interval_e, CS_Subevent_Interval, 625us);
        ENUM_W_MIN_MAX(HCI_CS_Test_Max_Num_Subevents_E, Max_Num_Subevents, 1, 20);
        ENUM_W_UNITS(HCI_Transmit_Power_Level_e, Transmit_Power_Level, dBm);
        HCI_T_IP1_time_e T_IP1_time;
        HCI_T_IP2_time_e T_IP2_time;
        HCI_T_FCS_time_e T_FCS_time;
        HCI_T_PM_time_e T_PM_time;
        HCI_T_SW_time_e T_SW_time;
        HCI_Tone_Antenna_Config_e Tone_Antenna_Config;
        HCI_Companion_Signal_Enable_Test_e Companion_Signal_Enable;
        HCI_SNR_Control_Initiator_e SNR_Control_Initiator;
        HCI_SNR_Control_Reflector_e SNR_Control_Reflector;
        U8 DRBG_Nonce[2];
        INT_W_MIN_MAX(U8, Channel_Map_Repetition, 1, 0xFF);
        HCI_Override_Config_t Override_Config;
        U8 PRINT_DEC(Override_Parameters_Length);
        U8 Override_Parameters_Data[1];
    }
    STRU_E_VLA1(HCI_CS_Test_Atlanta_r04_t);
    '''
    I = bytearray()
    I += s.pack('<BHB', 1, 0x2095, 0)
    I += s.pack('<8B', main_mode, sub_mode, main_mode_rep, mode0_steps, 0, rtt_type, phy, antenna)
    I += bits24(subevent_len)
    I += s.pack('<HBb6B', subevent_interval, max_num_subevents, TX_power, T_IP1, T_IP2, T_FCS, T_PM, T_SW, Tone_Antenna_Config)
    I += s.pack('<BBBHB', Companion_Signal_Enable, SNRs[0], SNRs[1], Nonce, Channel_Map_Repetition)

    add_override_data(I, AAs, channels, use_3C, hat_shape, ch3c_jump, markers, fewer_chans, channel_map)

    I[3] = len(I) - 4

    R = copy.deepcopy(I)
    R[8] = 1 # Set Role

    return I, R


class CS_Tester(devTester):
    def __init__(self, c : testerConfig):
        # Use 0 as central, 1 as advertiser
        devTester.__init__(self, c)

        self.CS_sec_count = 0
        self.config_ids = dict()
        self.has_TX_SNR = True

        self.input_cbs.append(self.on_CS_config_complete)

    def on_CS_config_complete(self, com_idx, pkt):
        if not pkt.hdr_matches([4, 0x3e, 0x22, 0x2f, 0x00]):
            return

        id = pkt.bytes[7]
        if id in self.config_ids:
            self.config_ids[id] = self.config_ids[id] + 1
            self.log('Board {}: CS Config {} seen *{}'.format(com_idx, id, self.config_ids[id]))
        else:
            self.config_ids[id] = 1

    def on_CS_security(self, com_idx, pkt):
        if not pkt.hdr_matches([4, 0x3e, 4, 0x2e, 0, self.conn_handles[com_idx] & 0xFF, self.conn_handles[com_idx] >> 8]):
            return

        self.log('Board {}: CS Security enabled'.format(com_idx))
        self.CS_sec_count = self.CS_sec_count + 1

    def CS_setup(self, timeout = 2):
        self.send_parallel_cmd_seqs((CS_setup_cmds(self.conn_handles[0]), CS_setup_cmds(self.conn_handles[1])))
        time.sleep(0.5)
        self.send_cmd_sequence(0, cmd_seq(['01 89 20 00'])) # LE CS Read Local Supported Capabilities
        self.send_cmd_sequence(1, cmd_seq(['01 89 20 00'])) # LE CS Read Local Supported Capabilities
        self.send_cmd_sequence(0, cmd_seq(['01 8A 20 02 ' + self.handle_bytes(0)])) # LE CS Read Remote Supported Capabilities
        time.sleep(0.5)
        self.send_cmd_sequence(1, cmd_seq(['01 8A 20 02 ' + self.handle_bytes(1)])) # LE CS Read Remote Supported Capabilities
        self.read_all()
        time.sleep(0.5)

        self.CS_sec_count = 0
        self.input_cbs.append(self.on_CS_security)
        self.send_cmd_sequence(0, cmd_seq(['01 8C 20 02 ' + self.handle_bytes(0)])) # LE CS Security Enable
        start = time.time()
        while self.CS_sec_count < 2 and time.time() < start + timeout:
            self.read_all()

    def send_CS_create_config(self, id, main_mode=2, sub_mode=1, rtt_type=0, min_main=2, max_main=6, main_mode_rep=1, n_mode0=3, 
                                    use_3C=False, ChMap_repeat=1, hat_shape=True, ch3c_jump=3, sync_phy=1, timeout=2):
        self.send_cmd_sequence(0, CS_create_config(self.conn_handles[0], id, main_mode, sub_mode, rtt_type, min_main, max_main, main_mode_rep, 
                                                   n_mode0, use_3C, ChMap_repeat, hat_shape, ch3c_jump, sync_phy))
        start = time.time()
        while not (id in self.config_ids and self.config_ids[id] >= 2) and time.time() < start + timeout:
            self.read_all()

    def send_CS_Test_end(self):
        self.send_cmd_sequence(0, cmd_seq(['01 96 20 00']))
        self.send_cmd_sequence(1, cmd_seq(['01 96 20 00']))

    def CS_run_proc(self, id, main_mode=3, sub_mode=2, rtt_type=4, sync_phy=1, ACI=0, min_main=2, max_main=6, main_mode_rep=1, n_mode0=3, subevent_len=8000, 
                              use_3C=False, ChMap_repeat=1, hat_shape=True, ch3c_jump=7, proc_count=1,
                              proc_interval=20, acl_phy=1):
        self.send_CS_create_config(id, main_mode, sub_mode, rtt_type, min_main, max_main, main_mode_rep, n_mode0, use_3C, ChMap_repeat, hat_shape, ch3c_jump, sync_phy)
        self.send_cmd_sequence(0, CS_set_proc_params(self.conn_handles[0], id, self.has_TX_SNR, proc_count, ACI, acl_phy, subevent_len, proc_interval))
        self.send_cmd_sequence(0, CS_procedure_enable(self.conn_handles[0], id))
        self.read_timeout(8)

    def CS_stop_proc(self, id):
        self.send_cmd_sequence(0, CS_procedure_enable(self.conn_handles[0], id, False))
        self.read_timeout(1)

    def CS_remove_configs(self, ids):
        for id in ids:
            cmd = cmd_seq(['01 91 20 03 ' + hp(s.pack('<HB', self.conn_handles[0], id))])
            self.send_cmd_sequence(0, cmd)

    def run_CS_Test(self, rev='BT6', main_mode=2, sub_mode=0xFF, main_mode_rep=0, mode0_steps=3, rtt_type=0, phy=1, antenna=1, subevent_len=4000000, subevent_interval=0,
                  max_num_subevents=0, TX_power=0, T_IP1=145, T_IP2=145, T_FCS=150, T_PM=40, T_SW=0,Tone_Antenna_Config=0, Companion_Signal_Enable=0,
                  Nonce=0, Channel_Map_Repetition=1, channels=list(range(2,23)) + list(range(26,77)), AAs=(0xC9D07A33,0x2F6C3998), use_3C=False,
                  hat_shape=True, ch3c_jump=3, fewer_chans=False, markers=None, SNRs = (0xFF, 0xFF), channel_map = None):
        if rev == 'r03':
            Init, Refl = CS_Test_dec23(main_mode, sub_mode, main_mode_rep, mode0_steps, rtt_type, phy, antenna, subevent_len,
                                       subevent_interval, max_num_subevents, TX_power, T_IP1, T_IP2, T_FCS, T_PM, T_SW, Tone_Antenna_Config,
                                       Companion_Signal_Enable, Nonce, Channel_Map_Repetition, channels, AAs, use_3C, hat_shape, ch3c_jump, fewer_chans, markers)
        elif rev == 'BT6':
            Init, Refl = CS_Test_BT6(main_mode, sub_mode, main_mode_rep, mode0_steps, rtt_type, phy, antenna, subevent_len,
                                     subevent_interval, max_num_subevents, TX_power, T_IP1, T_IP2, T_FCS, T_PM, T_SW, Tone_Antenna_Config,
                                     Companion_Signal_Enable, Nonce, Channel_Map_Repetition, channels, AAs, SNRs, use_3C, hat_shape, ch3c_jump, fewer_chans, markers, channel_map)
        self.send_cmd_sequence(1, [bytes_to_ints(Refl)])
        self.send_cmd_sequence(0, [bytes_to_ints(Init)])
        time.sleep(5)
        self.send_CS_Test_end()


def device_connect(coms=(4, 6), bauds=(115200, 115200)):
    t = CS_Tester(testerConfig(com_IDs=coms, bauds=bauds, HCI_to_pcapng=True))
    t.connect()
    t.encrypt()

    return t


def test_mode(coms=(5, 6), bauds = (115200, 115200)):
    t = CS_Tester(testerConfig(com_IDs = coms, bauds = bauds))

    return t

if __name__ == '__main__':
    t = device_connect()
    t.disconnect()
