#ifndef TUPLE_GEN_H_
#define TUPLE_GEN_H_

/*
 Wenxian Li

 Cisco Systems, Inc.
*/

/*
Copyright (c) 2015-2015 Cisco Systems, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <vector>
#include <algorithm>
#include <map>
#include <string>
#include <iostream>
#include <fstream>
#include <string>
#include <queue>
#include "common/c_common.h"
#include <bitset>
#include <yaml-cpp/yaml.h>
#include <mac_mapping.h>

#include <random>

class CTupleBase {
public:
       CTupleBase() {
           m_client_mac.inused = UNUSED;
       }
       uint32_t getClient() {
           return m_client_ip;
       }
       void setClient(uint32_t ip) {
           m_client_ip = ip;
       }
       uint32_t getClientId() {
           return m_client_idx;
       }
       void setClientId(uint32_t id) {
           m_client_idx = id;
       }
       
       uint32_t getServer(){
           return m_server_ip;
       }
       void setServer(uint32_t ip) {
           m_server_ip = ip;
       }
       uint32_t getServerId(){
           return m_server_idx;
       }
       void setServerId(uint32_t id) {
           m_server_idx = id;
       }
       uint16_t getServerPort() {
           return m_server_port;
       }
       void setServerPort(uint16_t port) {
           m_server_port = port;
       }
       uint16_t getClientPort() {
           return m_client_port;
       }
       void setClientPort(uint16_t port) {
           m_client_port = port;
       }
       mac_addr_align_t* getClientMac() {
           return &m_client_mac;
       }
       void setClientMac(mac_addr_align_t* mac_info) {
           if (mac_info != NULL) {
               memcpy(&m_client_mac, mac_info, sizeof(mac_addr_align_t));
               m_client_mac.inused = INUSED;
           } else {
               m_client_mac.inused = UNUSED;
           }
       }
       void setClientTuple(uint32_t ip,mac_addr_align_t*mac,uint16_t port) {
           setClient(ip);
           setClientMac(mac);
           setClientPort(port);
       }
       void setClientAll2(uint32_t id, uint32_t ip,uint16_t port) {
           setClientId(id);
           setClient(ip);
           setClientPort(port);
       }
 
       void setServerAll(uint32_t id, uint32_t ip) {
           setServerId(id);
           setServer(ip);
       }
       void getClientAll(uint32_t & id, uint32_t & ip, uint32_t & port) {
           id = getClientId();
           ip = getClient();
           port = getClientPort();
       }
       void getServerAll(uint32_t & id, uint32_t & ip) {
           id = getServerId();
           ip = getServer();
       }
private:
       uint32_t m_client_ip;
       uint32_t m_client_idx;
       uint32_t m_server_ip;
       uint32_t m_server_idx;
       mac_addr_align_t m_client_mac;
       uint16_t m_client_port;
       uint16_t m_server_port;
};



/*
 * Class that handle the client info
 */
#define MAX_CLIENTS 1000000
#define MAX_PORT (64000)
#define MIN_PORT (1024)
#define ILLEGAL_PORT (0)

#define PORT_FREE (0)
#define PORT_IN_USE (1)

/*FIXME*/
#define VLAN_SIZE (2)

#define FOREACH(vector) for(int i=0;i<vector.size();i++)


/* Client distribution */
typedef enum  {
    cdSEQ_DIST    = 0,
    cdRANDOM_DIST = 1,
    cdNORMAL_DIST = 2,
    cdMAX_DIST    = 3
} IP_DIST_t ;

/* For type 1, we generator port by maintaining a 64K bit array for each port.
 * In this case, we cannot support large number of clients due to memory exhausted. 
 *
 * So we develop a type 2 tuple generator. In this case, we only maintain a 16 bit
 * current port number for each client. To apply to type 2, it should meet:
 * number of clients > (longest_flow*Total_CPS)/64K
 *
 * TRex will decide which type to use automatically. It is transparent to users.
 * */

#define TYPE1 0
#define TYPE2 1
#define MAX_TYPE 3

class CIpInfoBase {
    public:
        virtual mac_addr_align_t* get_mac() { return NULL;}
        virtual void set_mac(mac_addr_align_t*){;}
        virtual uint16_t get_new_free_port() = 0;
        virtual void return_port(uint16_t a) = 0;
        virtual void generate_tuple(CTupleBase & tuple) = 0;
        virtual void return_all_ports() = 0;
        uint32_t get_ip() {
            return m_ip;
        }
        void set_ip(uint32_t ip) {
            m_ip = ip;
        }
        virtual ~CIpInfoBase() {}
    protected:
        uint32_t          m_ip;
};

//CClientInfo for large amount of clients support
class CIpInfoL : public CIpInfoBase {
 private:
    uint16_t m_curr_port;
 public:
    CIpInfoL() {
        m_curr_port = MIN_PORT;
    }
    uint16_t get_new_free_port() {
        if (m_curr_port>MAX_PORT) {
            m_curr_port = MIN_PORT;
        }
        return m_curr_port++;
    }

    void return_port(uint16_t a) {
    }

    void return_all_ports() {
        m_curr_port = MIN_PORT;
    }
};


class CIpInfo : public CIpInfoBase {
 private:
    std::bitset<MAX_PORT>  m_bitmap_port;
    uint16_t m_head_port;
    friend class CClientInfoUT;

 private:
    bool is_port_available(uint16_t port) {
        if (!is_port_legal(port)) {
            return PORT_IN_USE;
        }
        return (m_bitmap_port[port] == PORT_FREE);
    }

    /*
     * Return true if the port is legal
     *        false if the port is illegal.
     */
    bool is_port_legal(uint16_t port) {
        if (port>=MAX_PORT || port < MIN_PORT) {
                return false;
        }
        return true;
    }

    void m_head_port_set(uint16_t head) {
        if (!is_port_legal(head)) {
            return;
        }
        m_head_port = head;
    }

    // Try to find next free port
    void get_next_free_port_by_bit() {
        uint16_t cnt = 0;
        if (!is_port_legal(m_head_port)) {
            m_head_port = MIN_PORT;
        }
        while (true) {
            if (is_port_available(m_head_port)) {
                return;
            }
            cnt++;
            if (cnt>20) {
                /*FIXME: need to trigger some alarms?*/
                return;
            }
            m_head_port++;
            if (m_head_port>=MAX_PORT) {
                m_head_port = MIN_PORT;
            }
        }
    }


 public:
    CIpInfo() {
        m_head_port = MIN_PORT;
        m_bitmap_port.reset();
    }

    uint16_t get_new_free_port() {
        uint16_t r;

        get_next_free_port_by_bit();
        if (!is_port_available(m_head_port)) {
            m_head_port = MIN_PORT;
            return ILLEGAL_PORT;
        }

        m_bitmap_port[m_head_port] = PORT_IN_USE;
        r = m_head_port;
        m_head_port++;
        if (m_head_port>MAX_PORT) {
            m_head_port = MIN_PORT;
        }
        return r;
    }

    void return_port(uint16_t a) {
        assert(is_port_legal(a));
        assert(m_bitmap_port[a]==PORT_IN_USE);
        m_bitmap_port[a] = PORT_FREE;
    }

    void return_all_ports() {
        m_head_port = MIN_PORT;
        m_bitmap_port.reset();
    }
};

class CClientInfo : public CIpInfo {
public:
    CClientInfo (bool has_mac) {
        if (has_mac==true) {
            m_mac = new mac_addr_align_t();
        } else {
            m_mac = NULL;
        }
    }
    CClientInfo () {
        m_mac = NULL;
    }

    mac_addr_align_t* get_mac() {
        return m_mac;
    }
    void set_mac(mac_addr_align_t *mac) {
        memcpy(m_mac, mac, sizeof(mac_addr_align_t));
    }
    ~CClientInfo() {
        if (m_mac!=NULL){
            delete m_mac;
            m_mac=NULL;
        }
    }
    
    void generate_tuple(CTupleBase & tuple) {
        tuple.setClientTuple(m_ip, m_mac,
                           get_new_free_port());
    }
private:
    mac_addr_align_t *m_mac;
};

class CClientInfoL : public CIpInfoL {
public:
    CClientInfoL (bool has_mac) {
        if (has_mac==true) {
            m_mac = new mac_addr_align_t();
        } else {
            m_mac = NULL;
        }
    }
    CClientInfoL () {
        m_mac = NULL;
    }

    mac_addr_align_t* get_mac() {
        return m_mac;
    }

    void set_mac(mac_addr_align_t *mac) {
        memcpy(m_mac, mac, sizeof(mac_addr_align_t));
    }

    ~CClientInfoL() {
        if (m_mac!=NULL){
            delete m_mac;
            m_mac=NULL;
        }
    }
    
    void generate_tuple(CTupleBase & tuple) {
        tuple.setClientTuple(m_ip, m_mac,
                           get_new_free_port());
    }
private:
    mac_addr_align_t *m_mac;
};

class CServerInfo : public CIpInfo {
    void generate_tuple(CTupleBase & tuple) {
        tuple.setServer(m_ip);
    }
};

class CServerInfoL : public CIpInfoL {
    void generate_tuple(CTupleBase & tuple) {
        tuple.setServer(m_ip);
    }
};




class CIpPool {
    public:
       uint16_t GenerateOnePort(uint32_t idx) {
            CIpInfoBase* ip_info = m_ip_info[idx];
            uint16_t port;
            port = ip_info->get_new_free_port();
            
            //printf(" alloc extra  %x %d \n",c_ip,port);
            if (port==ILLEGAL_PORT) {
                m_port_allocation_error++;
            }
            m_active_alloc++;
            return (port);
        }

       bool is_valid_ip(uint32_t ip){
            CIpInfoBase* ip_front = m_ip_info.front();
            CIpInfoBase* ip_back  = m_ip_info.back();
            if ((ip>=ip_front->get_ip()) && 
                (ip<=ip_back->get_ip())) {
                return(true);
            }
            printf("invalid ip:%x, min_ip:%x, max_ip:%x, this:%p\n", 
                   ip, ip_front->get_ip(), 
                   ip_back->get_ip(),this);
            return(false);
        }

        uint32_t get_curr_ip() {
            return m_ip_info[m_cur_idx]->get_ip();
        }
        uint32_t get_ip(uint32_t idx) {
            return m_ip_info[idx]->get_ip();
        }
        CIpInfoBase* get_ip_info_by_idx(uint32_t idx) {
            return m_ip_info[idx];
        }

        void inc_cur_idx() {
            switch (m_dist) {
            case cdRANDOM_DIST: 
                m_cur_idx = get_random_idx();
                break;
            case cdSEQ_DIST :
            default:
                m_cur_idx++;
                if (m_cur_idx >= m_ip_info.size())
                    m_cur_idx = 0;
            }
        }
        //return a valid client idx in this pool
        uint32_t generate_ip() {
            uint32_t res_idx = m_cur_idx;
            inc_cur_idx();
            return res_idx;
        }

            

        void set_dist(IP_DIST_t dist) {
            if (dist>=cdMAX_DIST) {
                m_dist = cdSEQ_DIST;
            } else {
                m_dist = dist;
            }
        }
        void Delete() {
            FOREACH(m_ip_info) {
                delete m_ip_info[i];
            }
            m_ip_info.clear();
        }
        uint32_t get_total_ips() {
            return m_ip_info.size();
        }
        void return_all_ports() {
            FOREACH(m_ip_info) {
                m_ip_info[i]->return_all_ports();
            }
        }
        void FreePort(uint32_t id, uint16_t port) {
    //        assert(id<m_ip_info.size());
            m_active_alloc--;
            CIpInfoBase* client = m_ip_info[id];
            client->return_port(port);
        }
        
        mac_addr_align_t * get_curr_mac() {
            return m_ip_info[m_cur_idx]->get_mac();
        }
        mac_addr_align_t *get_mac(uint32_t idx) {
            return m_ip_info[idx]->get_mac();
        }
 
    public:
        std::vector<CIpInfoBase*> m_ip_info;
        IP_DIST_t  m_dist;
        uint32_t m_cur_idx;
        uint32_t m_active_alloc;
        uint32_t m_port_allocation_error;
        std::default_random_engine generator;
        std::uniform_int_distribution<int> *rand_dis;
        void CreateBase() {
            switch (m_dist) {
            case cdRANDOM_DIST:
                rand_dis = new std::uniform_int_distribution<int>
                    (0,get_total_ips()-1);
                break;
            default:
                break;
            }
            m_cur_idx = 0;
            m_active_alloc = 0;
            m_port_allocation_error = 0;
        }
        uint32_t get_random_idx() {
            uint32_t res =  (*rand_dis)(generator);
            return (res);
        }
        bool IsFreePortRequired(void){
            return(true);
        }


};

class CClientPool : public CIpPool {
public:

    uint32_t GenerateTuple(CTupleBase & tuple) {
        uint32_t idx = generate_ip();
        CIpInfoBase* ip_info = m_ip_info[idx];
        ip_info->generate_tuple(tuple);

        tuple.setClientId(idx);
        if (tuple.getClientPort()==ILLEGAL_PORT) {
            m_port_allocation_error++;
        }
        m_active_alloc++;
        return idx;
    }

    uint16_t get_tcp_aging() {
        return m_tcp_aging;
    }
    uint16_t get_udp_aging() {
        return m_udp_aging;
    }
    void Create(IP_DIST_t  dist_value,
                uint32_t min_ip,
                uint32_t max_ip,
                double l_flow,
                double t_cps,
                CFlowGenListMac* mac_info,
                bool has_mac_map, 
                uint16_t tcp_aging,
                uint16_t udp_aging); 

public: 
    uint16_t m_tcp_aging;
    uint16_t m_udp_aging;
};

class CServerPoolBase {
    public:

    virtual ~CServerPoolBase() {}

    virtual void GenerateTuple(CTupleBase& tuple) = 0;
    virtual uint16_t GenerateOnePort(uint32_t idx) = 0;
    virtual void Delete() = 0;
    virtual uint32_t get_total_ips()=0;
    virtual void Create(IP_DIST_t  dist_value,
               uint32_t min_ip,
               uint32_t max_ip,
               double l_flow,
               double t_cps) = 0; 
 
};

class CServerPoolSimple : public CServerPoolBase {
public:
    void Create(IP_DIST_t  dist_value,
               uint32_t min_ip,
               uint32_t max_ip,
               double l_flow,
               double t_cps) {
        m_max_server_ip = max_ip;
        m_min_server_ip = min_ip;
        m_cur_server_ip = min_ip;
    }
    void Delete() {
        return ;
    }
    void GenerateTuple(CTupleBase& tuple) {
        tuple.setServer(m_cur_server_ip);
        m_cur_server_ip ++;
        if (m_cur_server_ip > m_max_server_ip) {
            m_cur_server_ip = m_min_server_ip;
        }
    } 
    uint16_t GenerateOnePort(uint32_t idx) {
        // do nothing
        return 0;
    }
    uint32_t get_total_ips() {
        return (m_max_server_ip-m_min_server_ip+1);
    }
private:
    uint32_t m_max_server_ip;
    uint32_t m_min_server_ip;
    uint32_t m_cur_server_ip;
};

class CServerPool : public CServerPoolBase {
public:
    void GenerateTuple(CTupleBase & tuple) {
       uint32_t idx = gen->generate_ip();
       tuple.setServerAll(idx, gen->get_ip(idx));
    }
    uint16_t GenerateOnePort(uint32_t idx) {
        return gen->GenerateOnePort(idx);
    }
    void Create(IP_DIST_t  dist_value,
                uint32_t min_ip,
                uint32_t max_ip,
                double l_flow,
                double t_cps); 
 
    void Delete() {
        if (gen!=NULL) {
            gen->Delete();
            delete gen;
            gen=NULL;
        }
    }
    uint32_t get_total_ips() {
        return gen->m_ip_info.size();
    }
private: 
    CIpPool *gen;
};

/* generate for each template */
class CTupleGeneratorSmart {
public:
    /* return the active socket */
    uint32_t ActiveSockets(void){
        uint32_t total_active_alloc = 0;
        FOREACH(m_client_pool) {
            total_active_alloc += m_client_pool[i]->m_active_alloc;
        }
        return (total_active_alloc);
    }

    uint32_t getTotalClients(void){
        uint32_t total_clients = 0;
        FOREACH(m_client_pool) {
            total_clients += m_client_pool[i]->get_total_ips();
        }
        return (total_clients);
    }

    uint32_t getTotalServers(void){
        uint32_t total_servers = 0;
        FOREACH(m_server_pool) {
            total_servers += m_server_pool[i]->get_total_ips();
        }
        return total_servers;
    }

    uint32_t SocketsPerClient(void){
        return (MAX_PORT -MIN_PORT+1);
    }

    uint32_t MaxSockets(void){
        return (SocketsPerClient() * getTotalClients());
    }

    
    void FreePort(uint8_t pool_idx, uint32_t id, uint16_t port) {
        get_client_pool(pool_idx)->FreePort(id, port);
    }
        
    bool IsFreePortRequired(uint8_t pool_idx){
        return(get_client_pool(pool_idx)->IsFreePortRequired());
    }
    uint16_t get_tcp_aging(uint8_t pool_idx) {
        return (get_client_pool(pool_idx)->get_tcp_aging());
    }
    uint16_t get_udp_aging(uint8_t pool_idx) {
        return (get_client_pool(pool_idx)->get_udp_aging());
    }

public:
    CTupleGeneratorSmart(){
        m_was_init=false;
        m_has_mac_mapping = false;
    }
    bool Create(uint32_t _id,
            uint32_t thread_id, bool has_mac=false);

    void Delete();

    inline uint32_t GetThreadId(){
        return (  m_thread_id );
    }

    uint32_t getErrorAllocationCounter(){
        uint32_t total_alloc_error = 0;
        FOREACH(m_client_pool) {
            total_alloc_error += m_client_pool[i]->m_port_allocation_error;
        }
        return (total_alloc_error);
    }

    bool add_client_pool(IP_DIST_t  client_dist,
                         uint32_t min_client,
                         uint32_t max_client,
                         double l_flow,
                         double t_cps,
                         CFlowGenListMac* mac_info,
                         uint16_t tcp_aging,
                         uint16_t udp_aging);
    bool add_server_pool(IP_DIST_t  server_dist,
                         uint32_t min_server,
                         uint32_t max_server,
                         double l_flow,
                         double t_cps,
                         bool is_bundling);
    CClientPool* get_client_pool(uint8_t idx) {
        return m_client_pool[idx];
    }
    uint8_t get_client_pool_num() {
        return m_client_pool.size();
    }
    uint8_t get_server_pool_num() {
        return m_server_pool.size();
    }
    CServerPoolBase* get_server_pool(uint8_t idx) {
        return m_server_pool[idx];
    }
private:
    uint32_t m_id;
    uint32_t m_thread_id;
    std::vector<CClientPool*> m_client_pool;
    std::vector<CServerPoolBase*> m_server_pool;
    bool     m_was_init;
    bool     m_has_mac_mapping;
};

class CTupleTemplateGeneratorSmart {
public:
    /* simple tuple genertion for one low*/
    void GenerateTuple(CTupleBase & tuple){
        if (m_w==1) {
            /* new client each tuple generate */
            m_client_gen->GenerateTuple(tuple);
            m_server_gen->GenerateTuple(tuple);
            m_cache_client_ip = tuple.getClient();
            m_cache_client_idx = tuple.getClientId();
        }else{
            if (m_cnt==0) {
                m_client_gen->GenerateTuple(tuple);
                m_server_gen->GenerateTuple(tuple);
                m_cache_client_ip = tuple.getClient();
                m_cache_client_idx = tuple.getClientId();
                tuple.getServerAll(m_cache_server_idx, m_cache_server_ip);
            }else{
                tuple.setServerAll(m_cache_server_idx, 
                                   m_cache_server_ip);
                tuple.setClientAll2(m_cache_client_idx,
                                    m_cache_client_ip,
                                    m_client_gen->GenerateOnePort(m_cache_client_idx));
            }
            m_cnt++;
            if (m_cnt>=m_w) {
                m_cnt=0;
            }
        }
        if ( m_is_single_server ) {
            tuple.setServer(m_server_ip);
        }
    }

    uint16_t GenerateOneSourcePort(){
        return ( m_client_gen->GenerateOnePort(m_cache_client_idx) );
    }

    inline uint32_t GetThreadId(){
        return ( m_gen->GetThreadId() );
    }

public:

    bool Create( CTupleGeneratorSmart * gen,uint8_t c_pool,uint8_t s_pool){
        m_gen=gen;
        m_is_single_server=false;
        m_server_ip=0;
        SetW(1);
        m_client_gen = gen->get_client_pool(c_pool);
        m_server_gen = gen->get_server_pool(s_pool);
        return (true);
    }

    void Delete(){
    }
public:
    void SetW(uint16_t w){
        m_w=w;
        m_cnt=0;
    }

    uint16_t getW(){
        return (m_w);
    }


    void SetSingleServer(bool is_single, 
                         uint32_t server_ip,
                         uint32_t dual_port_index,
                         uint32_t dual_mask){
        m_is_single_server = is_single;
        m_server_ip = server_ip+dual_mask*dual_port_index;
    }
    bool IsSingleServer(){
        return (m_is_single_server);
    }

    CTupleGeneratorSmart * get_gen() {
        return m_gen;
    }
private:
    CTupleGeneratorSmart * m_gen;
    CClientPool          * m_client_gen;
    CServerPoolBase      * m_server_gen;
    uint16_t               m_w;
    uint16_t               m_cnt;
    uint32_t               m_server_ip;
    uint32_t               m_cache_client_ip;
    uint32_t               m_cache_client_idx;
    uint32_t               m_cache_server_ip;
    uint32_t               m_cache_server_idx;
    bool                   m_is_single_server;
};


/* YAML of generator */
#if 0
        -  client_distribution      : 'seq' - ( e.g c0,1,2,3,4  
                                      'random' - random from the pool 
                                      'normal' - need to give average and dev -- second phase      
                                        
        -  client_pool_mask         : 10.0.0.0-20.0.0.0
        -  server_pool_mask         : 70.0.0.0-70.0.20.0    
        -  number_of_clients_per_gb : 20
        -  dual_interface_mask      : 1.0.0.0  // each dual ports will add this to the pool of clients 
#endif

struct CTupleGenPoolYaml {
    IP_DIST_t       m_dist;
    uint32_t        m_ip_start;
    uint32_t        m_ip_end;
    uint32_t        m_number_of_clients_per_gb;
    uint32_t        m_min_clients;
    uint32_t        m_dual_interface_mask;
    uint16_t        m_tcp_aging_sec; /* 0 means there is no aging */
    uint16_t        m_udp_aging_sec;
    std::string     m_name;
    bool            m_is_bundling;
    public:
    uint32_t getTotalIps(void){
        return ( m_ip_end-m_ip_start+1);
    }
    uint32_t getDualMask() {
        return m_dual_interface_mask;
    }
    uint32_t get_ip_start() {
        return m_ip_start;
    }
    bool is_valid(uint32_t num_threads,bool is_plugins);
    void Dump(FILE *fd);
};
   

struct CTupleGenYamlInfo {
    std::vector<CTupleGenPoolYaml> m_client_pool;
    std::vector<CTupleGenPoolYaml> m_server_pool;
        
public:
    bool is_valid(uint32_t num_threads,bool is_plugins);
    uint8_t get_server_pool_id(std::string name){
        for (uint8_t i=0;i<m_server_pool.size();i++) {
            if (m_server_pool[i].m_name==name) 
                return i;
        }
        return 0;
    }

    uint8_t get_client_pool_id(std::string name){
        for (uint8_t i=0;i<m_client_pool.size();i++) {
            if (m_client_pool[i].m_name==name) 
                return i;
        }
        return 0;
    }
};


void operator >> (const YAML::Node& node, CTupleGenPoolYaml & fi) ;

void operator >> (const YAML::Node& node, CTupleGenYamlInfo & fi) ;


struct CIpPortion {
    uint32_t m_ip_start;
    uint32_t m_ip_end;
};
void split_ips(uint32_t thread_id,
                   uint32_t total_threads,
                   uint32_t dual_port_id,
                   CTupleGenPoolYaml& poolinfo,
                   CIpPortion & portion);



#endif //TUPLE_GEN_H_