import sys
sys.path.append('/nix/store/qfv2ayja1zgrfbiy1nrd9f0b1y759h60-python3-3.11.6-env/lib/python3.11/site-packages')
import logging
import bfrt_grpc.client as gc

logger = logging.getLogger('DigestListener')
logger.setLevel(logging.DEBUG)
log_handler = logging.FileHandler('digest_listener.log', mode='w')
log_handler.setLevel(logging.DEBUG)
logger.addHandler(log_handler)

HOST_IP = "localhost"
REGISTER_COUNTER = 'r_counter'
OTHER_REGISTERS = ['r_srcAddr', 'r_dstAddr']

def gc_connect():
    for bfrt_client_id in range(10):
        try:
            interface = gc.ClientInterface(
                grpc_addr=f"{HOST_IP}:50052",
                client_id=bfrt_client_id,
                device_id=0,
                num_tries=1,
            )
            break
        except Exception as e:
            logger.error("Failed to connect to gRPC server (client_id %d): %s", bfrt_client_id, str(e))
            print(f"Failed to connect to gRPC server (client_id {bfrt_client_id}): {e}")
            sys.exit(1)
    dev_tgt = gc.Target(device_id=0, pipe_id=0xffff)
    try:
        bfrt_info = interface.bfrt_info_get()
        if bfrt_client_id == 0:
            interface.bind_pipeline_config("cNdAmlight")
        logger.debug("Connected to gRPC server with client_id %d, p4_name: %s", bfrt_client_id, bfrt_info.p4_name_get())
        print(f"P4 program: {bfrt_info.p4_name_get()}")
        return interface, dev_tgt, bfrt_info
    except Exception as e:
        logger.error("Failed to get bfrt_info or bind pipeline: %s", str(e))
        print(f"Failed to initialize P4 program: {e}")
        sys.exit(1)

def clear_registers():
    interface, dev_tgt, bfrt_info = gc_connect()
    all_registers = [REGISTER_COUNTER] + OTHER_REGISTERS
    print("=" * 80)
    for register in all_registers:
        try:
            reg_name = f"SwitchIngress.{register}"
            table = bfrt_info.table_get(reg_name)
            table.entry_del(dev_tgt, [])
            print(f"Cleared register: {register}")
            logger.debug("Cleared register: %s", register)
        except Exception as e:
            print(f"Failed to clear register {register}: {e}")
            logger.error("Failed to clear register %s: %s", register, str(e))
    print("All registers cleared.")
    print("=" * 80)
    interface.tear_down_stream()

if __name__ == '__main__':
    clear_registers()