"""
Simple PTF test for super_vlan
"""

######### STANDARD MODULE IMPORTS ########
import unittest
import logging 
import grpc   
import pdb

######### PTF modules for BFRuntime Client Library APIs #######
import ptf
from ptf.testutils import *
from bfruntime_client_base_tests import BfRuntimeTest
import bfrt_grpc.bfruntime_pb2 as bfruntime_pb2
import bfrt_grpc.client as gc

########## Basic Initialization ############
class BaseProgramTest(BfRuntimeTest):
    def setUp(self):
        self.client_id = 0
        self.p4_name = "bloom_filter"
        self.dev      = 0
        self.dev_tgt  = gc.Target(self.dev, pipe_id=0xFFFF)
        
        print("\n")
        print("Test Setup")
        print("==========")

        BfRuntimeTest.setUp(self, self.client_id, self.p4_name)
        self.bfrt_info = self.interface.bfrt_info_get(self.p4_name)

        print("    Connected to Device: {}, Program: {}, ClientId: {}".format(
            self.dev, self.p4_name, self.client_id))

        # Since this class is not a test per se, we can use the setup method
        # for common setup. For example, we can have our tables and annotations
        # ready

        self.ipv4_host = self.bfrt_info.table_get("MyIngress.ipv4_host")
        self.ipv4_host.info.key_field_annotation_add(
            "hdr.ipv4.dst_addr", "ipv4")

        self.tables = [ self.ipv4_host ]

        # Create a list of all ports available on the device
        self.swports = []
        for (device, port, ifname) in ptf.config['interfaces']:
            self.swports.append(port)
        self.swports.sort()

        # Optional, but highly recommended
        self.cleanUp()

    # Use Cleanup Method to clear the tables before and after the test starts
    # (the latter is done as a part of tearDown()
    def cleanUp(self):
        print("\n")
        print("Table Cleanup:")
        print("==============")

        try:
            for t in self.tables:
                print("  Clearing Table {}".format(t.info.name_get()))
                keys = []
                for (d, k) in t.entry_get(self.dev_tgt):
                    if k is not None:
                        keys.append(k)
                t.entry_del(self.dev_tgt, keys)
                # Not all tables support default entry
                try:
                    t.default_entry_reset(self.dev_tgt)
                except:
                    pass
        except Exception as e:
            print("Error cleaning up: {}".format(e))

    # Use tearDown() method to return the DUT to the initial state by cleaning
    # all the configuration and clearing up the connection
    def tearDown(self):
        print("\n")
        print("Test TearDown:")
        print("==============")

        self.cleanUp()
        
        # Call the Parent tearDown
        BfRuntimeTest.tearDown(self)

#
# Individual tests can now be subclassed from BaseProgramTest
#

############################################################################
################# I N D I V I D U A L    T E S T S #########################
############################################################################

class insertBloomF(BaseProgramTest):
    def runTest(self):
        ingress_port = test_param_get("ingress_port",  0)
        ipv4_dst     = test_param_get("ipv4_dst",      "192.168.1.2")
        dst_addr     = test_param_get("dst_addr", "00:00:02:00:00:02")
        egress_port  = test_param_get("egress_port",   64)
        ipv4_src     = test_param_get("ipv4_dst",      "10.11.12.13")
        
        print("\n")
        print("Test Run")
        print("========")

        #
        # Program an entry in ipv4_host: ipv4_dst --> set_nexthop(nexthop_id)
        #
            
        #
        # Send a test packet
        #
        print("  Sending packet with IPv4 DST ADDR={} into port {}"
              .format(ipv4_dst, ingress_port))
        pkt = simple_tcp_packet(eth_dst=dst_addr,
                                eth_src='00:00:02:00:00:01',
                                ip_dst=ipv4_dst,
                                ip_src=ipv4_src,
                                ip_id=101,
                                ip_ttl=64,
                                ip_ihl=5,
                                tcp_sport=1234,
                                tcp_dport=123,
                                tcp_flags="S")
        send_packet(self, ingress_port, pkt)

        #
        # Wait for the egress packet and verify it
        #
        print("  Expecting the packet to be forwarded to port {}"
              .format(egress_port))
        verify_packet(self, pkt, egress_port)
        print("  Packet received of port {}".format(egress_port))

class hitBloomF(BaseProgramTest):
    def runTest(self):
        ingress_port = test_param_get("ingress_port",  0)
        ipv4_dst     = test_param_get("ipv4_dst",      "192.168.1.2")
        dst_addr     = test_param_get("dst_addr", "00:00:02:00:00:02")
        egress_port  = test_param_get("egress_port",   64)
        ipv4_src     = test_param_get("ipv4_dst",      "10.11.12.13")
        
        print("\n")
        print("Test Run")
        print("========")

        #
        # Program an entry in ipv4_host: ipv4_dst --> set_nexthop(nexthop_id)
        #

        #
        # Send a test packet
        #
        print("  Sending packet with IPv4 DST ADDR={} into port {}"
              .format(ipv4_dst, ingress_port))
        pkt = simple_tcp_packet(eth_dst=dst_addr,
                                eth_src='00:00:02:00:00:01',
                                ip_dst=ipv4_dst,
                                ip_src=ipv4_src,
                                ip_id=101,
                                ip_ttl=64,
                                ip_ihl=5,
                                tcp_sport=1234,
                                tcp_dport=123,
                                tcp_flags="FA")
        send_packet(self, ingress_port, pkt)

        #
        # Wait for the egress packet and verify it
        #
        print("  Expecting the packet to be forwarded to port {}"
              .format(egress_port))
        print("  Packet received of port {}".format(egress_port))


class NhitBloomF(BaseProgramTest):
    def runTest(self):
        ingress_port = test_param_get("ingress_port",  0)
        ipv4_dst     = test_param_get("ipv4_dst",      "192.168.1.2")
        dst_addr     = test_param_get("dst_addr", "00:00:02:00:00:02")
        egress_port  = test_param_get("egress_port",   64)
        ipv4_src     = test_param_get("ipv4_dst",      "101.11.12.13")
        
        print("\n")
        print("Test Run")
        print("========")

        #
        # Program an entry in ipv4_host: ipv4_dst --> set_nexthop(nexthop_id)
        #

        #
        # Send a test packet
        #
        print("  Sending packet with IPv4 DST ADDR={} into port {}"
              .format(ipv4_dst, ingress_port))
        pkt = simple_tcp_packet(eth_dst=dst_addr,
                                eth_src='00:00:02:00:00:01',
                                ip_dst=ipv4_dst,
                                ip_src=ipv4_src,
                                ip_id=101,
                                ip_ttl=64,
                                ip_ihl=5,
                                tcp_sport=12,
                                tcp_dport=123,
                                tcp_flags="FA")
        send_packet(self, ingress_port, pkt)

        #
        # Wait for the egress packet and verify it
        #
        print("  Expecting the packet to be forwarded to port {}"
              .format(egress_port)) 
        print("  Packet received of port {}".format(egress_port))