import sys
import time
import json
import argparse
import logging

from prov_p4.bmv2_controller import BMV2Controller


class ProvCollector:
    """This class collects the provenance from the data plane and sends it to the ProvManager for processing"""

    def __init__(
        self, reg_prov_details, cp_ip, thrift_port, switch_name, pkt_type, index_map
    ) -> None:
        self.reg_prov_details = reg_prov_details
        self.cp_ip = cp_ip
        self.dp_port = thrift_port
        self.controller = BMV2Controller(self.cp_ip, self.dp_port)
        self.agent_id = self.controller.get_switch_info()
        self.switch_name = switch_name
        self.pkt_type = pkt_type
        self.index_map = index_map
        self.last_count_index = {"packet_in": 0, "packet_out": 0}
        self.last_count_table_rule = 0
        for prov_reg_name in self.reg_prov_details:
            self.last_count_index[prov_reg_name] = 0
    
    def start(self): 
        POLL_TIME = 2  # In seconds
        try:
            while True:
                self._read_prov_regs()
                time.sleep(POLL_TIME)
        except KeyboardInterrupt:
            print("Closing Provenance Collector")

    def _read_prov_regs(self):
        for _reg_name in self.reg_prov_details:
            reg_name = _reg_name.replace(".")[1]
            if (
                reg_name == "packet_in"
                or (reg_name == "table_rule")
                or (reg_name == "packet_out")
            ):
                continue
            prov_reg_name = f"prov_{reg_name}"
            try:
                register_data = self.controller.read_register(prov_reg_name)
            except Exception as e:
                print(
                    f"[ERROR] Unable to read prov register (name - {prov_reg_name})...exiting"
                )
                logger.error(
                    f"Unable to read prov register (name - {prov_reg_name})...exiting",
                    e,
                )
                sys.exit(1)
            # Send Prov Agent event first
            record_agent(self.switch_name, self.agent_id, "Switch", time.time())
            prov_width_list = [
                ["index", self.reg_prov_details[reg_name]["index_size"]],
                ["value", self.reg_prov_details[reg_name]["value_size"]],
                ["activity", ACTIVITY_SIZE],
                ["version", VERSION_SIZE],
            ]
            # 

def main():
    """This script takes the provenance details, extracts the DP provenance, parses it
    and sends it to the RabbitMQ for processing using SPADE"""
    parser = argparse.ArgumentParser(
        description="This script requires the provenance detail file and thrift port as input."
    )
    parser.add_argument("--prov-details", type=str, help="Provenance details JSON")
    parser.add_argument("--switch-name", type=str, help="Switch Name (s0,s1...)")
    parser.add_argument(
        "--pkt-type", type=str, help="PacketIn/PacketOut field types JSON configuration"
    )
    parser.add_argument("--index-map", type=str, help="Index map JSON configuration")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "-f", "--thrift-port-file", type=str, help="Path to the thrift port file."
    )
    group.add_argument("-p", "--thrift-port", type=int, help="Thrift port directly.")
    args = parser.parse_args()

    if args.thrift_port_file:
        with open(args.thrift_port_file, "r") as file:
            line = file.readline().strip()
            thrift_port = int(line)
    elif args.thrift_port is not None:
        thrift_port = args.thrift_port
    cp_ip = "localhost"
    args = parser.parse_args()
    # Read all the configuration files - provenance details, index map, and packet annotations

    with open(args.prov_details, "r") as file:
        reg_prov_details = json.load(file)
    with open(args.index_map, "r") as file:
        index_map = json.load(file)
    with open(args.pkt_type, "r") as file:
        pkt_type = json.load(file)
    print(reg_prov_details)
    print(index_map)
    print(pkt_type)
    collector = ProvCollector(
        reg_prov_details, cp_ip, thrift_port, args.switch_name, pkt_type, index_map
    )
    collector.start()


if __name__ == "__main__":
    main()
