#!/router/bin/python

import pprint
import yaml
import os
# import bisect

class CTRexYaml(object):
    """
    This class functions as a YAML generator according to TRex YAML format.

    CTRexYaml is compatible with both Python 2 and Python 3.
    """
    YAML_TEMPLATE = [{'cap_info': [],
                      'duration': 10.0,
                      'generator': {'clients_end': '16.0.1.255',
                                    'clients_per_gb': 201,
                                    'clients_start': '16.0.0.1',
                                    'distribution': 'seq',
                                    'dual_port_mask': '1.0.0.0',
                                    'min_clients': 101,
                                    'servers_end': '48.0.0.255',
                                    'servers_start': '48.0.0.1',
                                    'tcp_aging': 1,
                                    'udp_aging': 1},
                      'mac' :  [0x00,0x00,0x00,0x01,0x00,0x00]}]
    PCAP_TEMPLATE = {'cps': 1.0,
                    'ipg': 10000,
                    'name': '',
                    'rtt': 10000,
                    'w': 1}

    def __init__ (self, trex_files_path):
        """
        The initialization of this class creates a CTRexYaml object with **empty** 'cap-info', and with default client-server configuration.

        Use class methods to add and assign pcap files and export the data to a YAML file.

        :parameters:  
            trex_files_path : str
                a path (on TRex server side) for the pcap files using which TRex can access it.

        """
        self.yaml_obj        = list(CTRexYaml.YAML_TEMPLATE)
        self.empty_cap       = True
        self.file_list       = []
        self.yaml_dumped     = False
        self.trex_files_path = trex_files_path

    def add_pcap_file (self, local_pcap_path):
        """
        Adds a .pcap file with recorded traffic to the yaml object by linking the file with 'cap-info' template key fields.
                
        :parameters:  
            local_pcap_path : str
                a path (on client side) for the pcap file to be added.

        :return: 
            + The index of the inserted item (as int) if item added successfully
            + -1 if pcap file already exists in 'cap_info'.

        """
        new_pcap = dict(CTRexYaml.PCAP_TEMPLATE)
        new_pcap['name'] = self.trex_files_path + os.path.basename(local_pcap_path)
        if self.get_pcap_idx(new_pcap['name']) != -1:
            # pcap already exists in 'cap_info'
            return -1
        else:
            self.yaml_obj[0]['cap_info'].append(new_pcap)
            if self.empty_cap:
                self.empty_cap = False
            self.file_list.append(local_pcap_path)
            return ( len(self.yaml_obj[0]['cap_info']) - 1)


    def get_pcap_idx (self, pcap_name):
        """
        Checks if a certain .pcap file has been added into the yaml object.
                
        :parameters:  
            pcap_name : str
                the name of the pcap file to be searched

        :return: 
            + The index of the pcap file (as int) if exists
            + -1 if not exists.

        """
        comp_pcap = pcap_name if pcap_name.startswith(self.trex_files_path) else (self.trex_files_path + pcap_name)
        for idx, pcap in enumerate(self.yaml_obj[0]['cap_info']):
            print (pcap['name'] == comp_pcap)
            if pcap['name'] == comp_pcap:
                return idx
        # pcap file wasn't found
        return -1

    def dump_as_python_obj (self):
        """
        dumps with nice indentation the pythonic format (dictionaries and lists) of the currently built yaml object.
                
        :parameters:  
            None

        :return: 
            None

        """
        pprint.pprint(self.yaml_obj)

    def dump(self):
        """
        dumps with nice indentation the YAML format of the currently built yaml object.
                
        :parameters:  
            None

        :return:
            None

        """
        print (yaml.safe_dump(self.yaml_obj, default_flow_style = False))

    def to_yaml(self, filename):
        """
        Exports to YAML file the built configuration into an actual YAML file.
                
        :parameters:  
            filename : str
                a path (on client side, including filename) to store the generated yaml file.

        :return: 
            None

        :raises:
            + :exc:`ValueError`, in case no pcap files has been added to the object.
            + :exc:`EnvironmentError`, in case of any IO error of writing to the files or OSError when trying to open it for writing.
        
        """
        if self.empty_cap:
            raise ValueError("No .pcap file has been assigned to yaml object. Must add at least one")
        else:
            try:
                with open(filename, 'w') as yaml_file:
                    yaml_file.write( yaml.safe_dump(self.yaml_obj, default_flow_style = False) )
                self.yaml_dumped = True
                self.file_list.append(filename)
            except EnvironmentError as inst:
                raise 

    def set_cap_info_param (self, param, value, seq):
        """
        Set cap-info parameters' value of a specific pcap file.
                
        :parameters:  
            param : str
                the name of the parameters to be set.
            value : int/float
                the desired value to be set to `param` key.
            seq : int
                an index to the relevant caps array to be changed (index supplied when adding new pcap file, see :func:`add_pcap_file`).

        :return: 
            **True** on success

        :raises:
            :exc:`IndexError`, in case an out-of range index was given.
        
        """
        try:
            self.yaml_obj[0]['cap_info'][seq][param] = value

            return True
        except IndexError:
            return False

    def set_generator_param (self, param, value):
        """
        Set generator parameters' value of the yaml object.
                
        :parameters:  
            param : str
                the name of the parameters to be set.
            value : int/float/str
                the desired value to be set to `param` key.

        :return: 
            None
        
        """
        self.yaml_obj[0]['generator'][param] = value

    def get_file_list(self):
        """
        Returns a list of all files related to the YAML object, including the YAML filename itself.

        .. tip:: This method is especially useful for listing all the files that should be pushed to TRex server as part of the same yaml selection.
                
        :parameters:  
            None

        :return: 
            a list of filepaths, each is a local client-machine file path.
        
        """
        if not self.yaml_dumped:
            print ("WARNING: .yaml file wasn't dumped yet. Files list contains only .pcap files")
        return self.file_list



if __name__ == "__main__":
    pass