# vim: syntax=py
# -*- mode: python -*-
# SPDX-License-Identifier: Apache-2.0
# Copyright(c) 2019 Intel Corporation

# for errnos
import errno
from conf.utils import *


# ====================================================
#       Parameters
# ====================================================


conf_file = get_env('CONF_FILE', 'conf/spgwu.json')
conf = get_json_conf(conf_file, False)

# Maximum number of flows to manage ip frags for re-assembly
max_ip_defrag_flows = None
try:
    max_ip_defrag_flows = int(conf["max_ip_defrag_flows"])
except ValueError:
    print('Invalid value for max_ip_defrag_flows. Not installing IPDefrag module.')
except KeyError:
    print('max_ip_defrag_flows value not set. Not installing IPDefrag module.')

# ====================================================
#       Port Helpers
# ====================================================


MAX_GATES = 8192
dpdk_ports = {}


def scan_dpdk_ports():
    idx = 0
    while True:
        try:
            intf = PMDPort(name="Port {}".format(idx), port_id=idx)
            if intf:
                # Need to declare mac so that we don't lose key during destroy_port
                mac = intf.mac_addr
                dpdk_ports[mac] = idx
                bess.destroy_port(intf.name)
        except bess.Error as e:
            if e.code == errno.ENODEV:
                break
            else:
                raise
        idx += 1
        # RTE_MAX_ETHPORTS is 32 and we need 2 for vdevs
        if idx == 30:
          break
    return True if dpdk_ports else False


class Port:
    def __init__(self, name):
        self.name = name
        self.wid = None
        self.fpi = None
        self.fpo = None
        self.bpf = None
        self.bpfgate = 0
        self.routes_table = None

    def bpf_gate(self):
        if self.bpfgate < MAX_GATES - 2:
            self.bpfgate += 1
            return self.bpfgate
        else:
            raise Exception('Port {}: Out of BPF gates to allocate'.format(self.name))

    def detect_mode(self):
        # default case
        mode = "unselected"

        try:
            peer_by_interface(self.name)
            mode = "dpdk"
        except:
            mode = "linux"
        return mode

    def init_fastpath(self, **kwargs):
        # Initialize PMDPort and RX/TX modules
        name = self.name
        fast = PMDPort(name="{}Fast".format(name), **kwargs)
        self.fpi = __bess_module__("{}FastPI".format(name), 'PortInc', port=fast.name)
        self.fpo = __bess_module__("{}FastPO".format(name), 'PortOut', port=fast.name)

        # Initialize BPF to classify incoming traffic to go to kernel and/or pipeline
        self.bpf = __bess_module__("{}FastBPF".format(name), 'BPF')
        self.bpf.clear()

        # Reassemble IP4 fragments (if needed)
        defrag = __bess_module__("{}IP4Defrag".format(name), 'IPDefrag', num_flows=max_ip_defrag_flows, numa=0)
        # Default drop when no matches
        self.fpi -> defrag:1 -> self.bpf:0 -> Sink()
        defrag:0 -> Sink()

        # Initialize route module
        self.rtr = __bess_module__("{}Routes".format(name), 'IPLookup')

        # Default route goes to Sink
        self.rtr.add(prefix='0.0.0.0', prefix_len=0, gate=MAX_GATES-1)
        self.rtr:(MAX_GATES-1) -> Sink()

        # Attach fastpath to worker's root TC
        self.fpi.attach_task(wid=self.wid)

    def setup_port(self, idx, workers):
        # Pick the worker handling this port
        self.wid = idx % workers

        name = self.name
        wid = self.wid
        print('Setting up port {} on worker {}'.format(name,wid))

        # Detect the mode of this interface - DPDK/AF_XDP/AF_PACKET
        mode = self.detect_mode()

        if mode == 'linux':
            try:
                # Initialize kernel fastpath.
                # AF_XDP requires that num_rx_qs == num_tx_qs
                kwargs = {"vdev" : "net_af_xdp{},iface={},start_queue=0,queue_count={}"
                          .format(idx, name, workers), "num_out_q": workers, "num_inc_q": workers}
                self.init_fastpath(**kwargs)
            except:
                print('Failed to create AF_XDP socket for {}. Creating AF_PACKET socket instead.'.format(name))
                kwargs = {"vdev" : "net_af_packet{},iface={},qpairs={}".format(idx, name, workers), "num_out_q": workers}
                self.init_fastpath(**kwargs)

        elif mode == 'dpdk':
            # if port list is empty, scan for dpdk_ports first
            if not dpdk_ports and scan_dpdk_ports() == False:
                print('Registered dpdk ports do not exist.')
                sys.exit()
            # Initialize DPDK fastpath
            fidx = dpdk_ports.get(mac_by_interface(name))
            if fidx is None:
                raise Exception('Registered port for {} not detected!'.format(name))
            kwargs = {"port_id" : fidx, "num_out_q": workers}
            self.init_fastpath( **kwargs)

            # Initialize kernel slowpath port and RX/TX modules
            try:
                peer = peer_by_interface(name)
                vdev = "net_af_packet{},iface={}".format(idx, peer)
                slow = PMDPort(name="{}Slow".format(name), vdev=vdev)
                spi = __bess_module__("{}SlowPI".format(name), 'PortInc', port=slow.name)
                spo = __bess_module__("{}SlowPO".format(name), 'PortOut', port=slow.name)
                qspo = __bess_module__("{}QSlowPO".format(name), 'Queue')

                # host_ip_filter: tcpdump -i foo 'dst host 198.19.0.1 or 198.18.0.1' -d
                # Should always be set to lowest priority
                HostGate = MAX_GATES - 1
                ips = ips_by_interface(name)
                host_ip_filter = {"priority": -HostGate, "filter": "dst host "
                                + " or ".join(str(x) for x in ips), "gate": HostGate}

                self.bpf.add(filters=[host_ip_filter])

                # Direct control traffic from DPDK to kernel
                self.bpf:HostGate -> qspo -> spo

                # Direct control traffic from kernel to DPDK
                spi -> self.fpo

                # Direct fast path traffic to Frag module
                merge = __bess_module__("{}Merge".format(name), 'Merge')
                frag = __bess_module__("{}IP4Frag".format(name), 'IPFrag')
                merge -> frag:1 -> self.fpo
                frag:0 -> Sink()

                tc = 'slow{}'.format(wid)
                try:
                    bess.add_tc(tc, policy='round_robin', wid=wid)
                except Exception as e:
                    if e.errmsg == "Name '{}' already exists".format(tc):
                        pass
                    else:
                        raise e
                # Limit scheduling slow path RX/TX to 1000 times/second each
                for mod in spi, qspo:
                    bess.add_tc(mod.name,
                            parent=tc,
                            policy='rate_limit',
                            resource='count',
                            limit={'count': 1000})
                    mod.attach_task(mod.name)
            except Exception as e:
                print('Mirror veth interface: {} misconfigured: {}'.format(name, e))
        else:
            raise Exception('Invalid mode selected.')


# ====================================================
#       Validate
# ====================================================


# CIDRs
enb_cidr = conf["enb_cidr"]
ue_cidr = conf["ue_cidr"]

cidrs = [enb_cidr, ue_cidr]
for cidr in cidrs:
    if validate_cidr(cidr) is False:
        exit(1, 'Invalid {}'.format(cidr))

# ====================================================
#       Core Setup
# ====================================================


# Initialize workers
workers = 1
try:
    workers = int(conf["workers"])
except ValueError:
    print('Invalid workers value! Re-setting # of workers to 1.')

cores = get_process_affinity()
for wid in xrange(workers):
    bess.add_worker(wid=wid, core=int(cores[wid % len(cores)]))

# ====================================================
#       Port Setup
# ====================================================


interfaces = ["s1u", "sgi"]
ports = {}
for idx, interface in enumerate(interfaces):
    port = Port(conf[interface]["ifname"])
    if port.name in ports:
        continue
    port.setup_port(idx, workers)
    ports[port.name] = port

s1u_ifname = conf["s1u"]["ifname"]
sgi_ifname = conf["sgi"]["ifname"]

# ====================================================
#       Downlink Pipeline
# ====================================================

s1uRoutes = ports[s1u_ifname].rtr

# Maximum number of sessions to manage
try:
    max_sessions = int(conf["max_sessions"])
except ValueError:
    print('Invalid max_sessions value!')

# Setting filter to detect UE subnet
sgiFastBPF = ports[sgi_ifname].bpf
UEGate = ports[sgi_ifname].bpf_gate()
ue_filter = {"priority": -UEGate,
             "filter": "ip dst net {}".format(ue_cidr), "gate": UEGate}
sgiFastBPF.add(filters=[ue_filter])

sgiFastBPF:UEGate \
    -> EtherTrim::GenericDecap(bytes=14) \
    -> GTPUEncap::GtpuEncap(s1u_sgw_ip=ip2long(ips_by_interface(s1u_ifname)[0]), num_subscribers=max_sessions):1 \
    -> S1UEtherAdd::GenericEncap(fields=[
        {'size': 6, 'value': {'value_int': 0x0}},
        {'size': 6, 'value': {'value_int': mac2hex(mac_by_interface(s1u_ifname))}},
        {'size': 2, 'value': {'value_int': 0x0800}}]) \
    -> OuterUDPCsum::L4Checksum() \
    -> OuterIPCsum::IPChecksum() \
    -> s1uRoutes

# Drop unknown packets
GTPUEncap:0 -> Sink()

# ====================================================
#       Uplink Pipeline
# ====================================================


# Setting filter to detect gtpu traffic
# src net 11.1.1.0 mask 255.255.255.0           # check eNB subnet
# and dst host 11.1.1.1                         # check S/PGWU IP
# and udp dst port 2152                         # check GTPU port
# and (udp[28:4] & 0xffffff00) = 0x10000000     # check UE subnet
s1uFastBPF = ports[s1u_ifname].bpf
check_enb_subnet = "src net {} ".format(enb_cidr)
check_spgwu_ip = " and dst host " + \
    " or ".join(str(x) for x in ips_by_interface(s1u_ifname))
check_gtpu_port = " and udp dst port 2152"
check_ue_subnet = " and (udp[28:4] & 0x{}) = 0x{}".format(
    cidr2mask(ue_cidr), ip2hex(cidr2block(ue_cidr)[0]))
check_gtpu_msg_echo = " and udp[9] = 0x1"

GTPUEchoGate = ports[s1u_ifname].bpf_gate()
uplink_echo_filter = {"priority": -GTPUEchoGate, "filter": check_enb_subnet +
                      check_spgwu_ip + check_gtpu_port +
                      check_gtpu_msg_echo, "gate": GTPUEchoGate}
s1uFastBPF.add(filters=[uplink_echo_filter])

GTPUGate = ports[s1u_ifname].bpf_gate()
uplink_filter = {"priority": -GTPUGate, "filter": check_enb_subnet +
               check_spgwu_ip + check_gtpu_port, "gate": GTPUGate}
s1uFastBPF.add(filters=[uplink_filter])

sgiRoutes = ports[sgi_ifname].rtr

s1uFastBPF:GTPUGate \
    -> EtherDecapTrim::GenericDecap(bytes=14) -> GTPUDecap::GtpuDecap(ename="GTPUEncap"):1 \
    -> SGIEtherAdd::GenericEncap(fields=[
        {'size': 6, 'value': {'value_int': 0x0}},
        {'size': 6, 'value': {'value_int': mac2hex(mac_by_interface(sgi_ifname))}},
        {'size': 2, 'value': {'value_int': 0x0800}}]) \
    -> sgiRoutes

s1uFastBPF:GTPUEchoGate \
	-> GTPUEcho::GtpuEcho(s1u_sgw_ip=ip2long(ips_by_interface(s1u_ifname)[0])):1 \
        -> EthSwap::MACSwap() \
        -> 1:OuterUDPCsum

# Drop unknown packets
GTPUEcho:0 -> Sink()
GTPUDecap:0 -> Sink()

# ====================================================
#       SIM_TEST
# ====================================================

import time

def sim_start_test():
    start_teid = 0xf0000000
    start_ue_ip = 0x10000001
    start_enb_ip = 0x0b010181
    NG4T_MAX_UE_RAN = 500000
    NG4T_MAX_ENB_RAN = 80

    for i in range(max_sessions):

        # NG4T-based formula to calculate enodeB IP address against a given UE IP address
        # il_trafficgen also uses the same scheme
        # See SimuCPEnbv4Teid(...) in ngic code for more details
        ue_of_ran = i % NG4T_MAX_UE_RAN
        ran = i / NG4T_MAX_UE_RAN
        enb_of_ran = ue_of_ran % NG4T_MAX_ENB_RAN
        enb_idx = ran * NG4T_MAX_ENB_RAN + enb_of_ran

        GTPUEncap.add(teid=start_teid+i, eteid=i+1, ueaddr=start_ue_ip+i, enodeb_ip=start_enb_ip+enb_idx)

def sim_end_test():
    start_ue_ip = 0x10000001

    for i in range(max_sessions):
        GTPUEncap.remove(ueaddr=start_ue_ip+i)

    GTPUEncap.show_records()

### Uncomment the following lines to test with il_trafficgen ###
# Start the test
#sim_start_test()
# Wait for 30 seconds before deleting the session info entries (optional)
#time.sleep(30)
#sim_end_test()

# Finally send SIGHUP to route_control daemon on reload
# TODO: behavior is unspecified if route_control.py pid is not found
route_control_pid = getpythonpid('route_control.py')
if route_control_pid:
    os.kill(route_control_pid, signal.SIGHUP)
