import sys
import time
import json
import argparse
import logging
import ipaddress

from provp4_api import (
    record_activity,
    record_API_call,
    record_agent,
    close_channel,
)
from bmv2_controller import BMV2Controller

# Static global variables for index
REG_NAME_INDEX = 0
ACTION_NAME_INDEX = 1
VALUE_ANNOTATION_INDEX = 2
INDEX_SIZE_INDEX = 3
VALUE_SIZE_INDEX = 4

# CONST SIZES # TODO - Make this come from a config file
ACTIVITY_SIZE = 8  # 8 bits for activity ID - Should be computed
# ACTIVITY_SIZE = 3
OPERATION_SIZE = 1  # 1 bit for READ/WRITE tag
VERSION_SIZE = 8
# VERSION_SIZE = 3

REGISTER_SIZE = 65536
# PKT_IN_TYPE = {"ipid": "hex", "src_ip": "ip_addr"}
PKT_IN_TYPE = {"src_mac": "mac_addr", "packet_id": "hex", "ether_type": "hex"}
index_map = {"MyIngress.prov_MyIngress_status_register": 0}

logger = logging.getLogger("dp-prov")


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, thrift_port, switch_name) -> None:
        self.reg_prov_details = reg_prov_details
        self.dp_ip = "localhost"
        self.dp_port = thrift_port
        self.cp_client = None
        self.controller = BMV2Controller("localhost", thrift_port)
        self.agent_id = self.controller.get_switch_info()
        self.switch_name = switch_name
        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 _parse_prov_test(self, reg_name, total_size, entry_string) -> dict:
        total_bit_width = sum([x[1] for x in total_size])
        logger.info(f"total bit width: {total_bit_width}")
        prov_string = format(entry_string, "b").zfill(total_bit_width + OPERATION_SIZE)
        logger.info(f"reg name: {reg_name} and prov_string: {prov_string}")
        start_index = 0
        entity_value_dict = {}
        operation_tag = prov_string[-1]
        prov_string = prov_string[:-1]  # Remove the operation bit
        new_size = []
        entity_value_type = self.reg_prov_details.get(reg_name).get("value_type")
        if operation_tag == "0":
            # READ operation. Therfore, swap the index and value sizes
            # total_size[0], total_size[1] = total_size[1], total_size[0]
            new_size.append(total_size[1])
            new_size.append(total_size[0])
            new_size.extend(total_size[2:])
            # logger.info(f"size: {new_size}")
            entity_value_dict["operation"] = "READ"
        else:
            new_size.append(total_size[0])
            new_size.append(total_size[1])
            new_size.extend(total_size[2:])
            # logger.info(f"size: {new_size}")
            entity_value_dict["operation"] = "WRITE"
        for entity_field in new_size:
            entity_type = entity_field[0]
            entity_size = entity_field[1]
            end_index = start_index + entity_size
            entity_value = prov_string[start_index:end_index]
            if entity_type == "value" and entity_value_type == "ip_addr":
                entity_value_converted = ".".join(
                    [
                        str((int(entity_value, 2) >> (8 * i)) & 255)
                        for i in range(3, -1, -1)
                    ]
                )
            elif entity_type == "value" and entity_value_type == "mac_addr":
                entity_value_converted = "{:012x}".format(int(entity_value, 2))
                entity_value_converted = ":".join(
                    [entity_value_converted[i : i + 2] for i in range(0, 12, 2)]
                )
            else:
                entity_value_converted = int(entity_value, 2)
            start_index = end_index
            entity_value_dict.update({entity_type: entity_value_converted})
            # logger.info(f"{entity_type}, {entity_value_converted}")
        # if entity_value_dict["operation"] == 0:
        #     logger.info(f"Value: {entity_value_dict['value']} - READ")
        # elif entity_value_dict["operation"] == 1:
        #     logger.info(f"Value: {entity_value_dict['value']} - WRITE")
        # else:
        #     logger.info("Other")
        logger.info(entity_value_dict)
        # Add the activity name instead of the activity id
        activity_id = str(entity_value_dict["activity"])
        if activity_id == "255":
            print("[WARN] 255 detected")
        else:
            entity_value_dict.update(
                {"activity": self.reg_prov_details[reg_name][activity_id]}
            )
        return entity_value_dict

    def _prov_wraparound(self):
        entries_read = 0
        for reg_name in self.reg_prov_details:
            if (
                reg_name == "packet_in"
                or (reg_name == "table_rule")
                or (reg_name == "packet_out")
            ):
                continue
            # TODO - make this dynamic, carefully check the P4 variable names
            prov_reg_name = f"MyIngress.prov_{reg_name}"
            # prov_reg_name = f"ingress.firewall_control.prov_{reg_name}"
            # prov_reg_name = f"prov_{reg_name}"
            # prov_reg_name = f"MyIngress.prov_{reg_name}"
            try:
                regs = 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)
            # Query the index array and save the last written index
            # prov_reg_indices = self.controller.read_register(
            #     "ingress.firewall_control.prov_index_reg"
            # )
            prov_reg_indices = self.controller.read_register("MyIngress.prov_index_reg")
            # prov_reg_indices = self.controller.read_register("prov_index_reg")
            end_idx = prov_reg_indices[index_map[prov_reg_name]]
            start_idx = self.last_count_index[reg_name]
            idx = start_idx

            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],
            ]

            while idx != end_idx:
                prov_entry = regs[idx]
                if prov_entry == 0:
                    print(f"Entry 0: idx: {idx} and end_idx:{end_idx}")
                    break
                # Parse the provenance record
                parsed_prov_values = self._parse_prov_test(
                    reg_name, prov_width_list, prov_entry
                )
                record_activity(
                    parsed_prov_values["activity"],
                    self.agent_id,
                    parsed_prov_values["version"],
                    "PacketProcessor",
                )
                converted_value = parsed_prov_values["value"]
                record_API_call(
                    reg_name,
                    self.agent_id,
                    parsed_prov_values["version"],
                    "register",
                    parsed_prov_values["operation"],
                    {
                        "index": parsed_prov_values["index"],
                        "value": converted_value,
                        "activity_name": parsed_prov_values["activity"],
                    },
                )
                logger.info(f"Register ({reg_name}) entity value: {parsed_prov_values}")
                entries_read += 1
                idx = (idx + 1) % REGISTER_SIZE
                # print(f"Entry: {prov_entry} and Index: {idx}")
            self.last_count_index[reg_name] = end_idx
            print(
                f"[INFO] {reg_name}: start={start_idx}, end={idx}, total_entries_read={entries_read}, {round((self.last_count_index[reg_name] / REGISTER_SIZE) * 100, 3)}"
            )
            # print(
            #     f"[INFO] Pulling Register ({reg_name}) provenance.. {round((self.last_count_index[reg_name] / REGISTER_SIZE) * 100, 3)}"
            # )

    def _read_prov_regs(self):
        for reg_name in self.reg_prov_details:
            if (
                reg_name == "packet_in"
                or (reg_name == "table_rule")
                or (reg_name == "packet_out")
            ):
                continue
            # TODO - make this dynamic, carefully check the P4 variable names
            # prov_reg_name = f"prov_{reg_name}"
            # prov_reg_name = f"ingress.firewall_control.prov_{reg_name}"
            prov_reg_name = f"MyIngress.prov_{reg_name}"
            try:
                regs = 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)
            count = 0
            # register_name = self.reg_prov_details[prov_reg_name][REG_NAME_INDEX]
            # get the activity name after parsing prov record
            # activity_name = self.reg_prov_details[prov_reg_name][ACTION_NAME_INDEX]
            # value_annotation = self.reg_prov_details[prov_reg_name][
            #     VALUE_ANNOTATION_INDEX
            # ]
            # activity_sent = False
            # Send the agent and activity once
            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],
            ]
            for prov_entry in regs[self.last_count_index[reg_name] :]:
                if prov_entry != 0:
                    count += 1
                    # Parse the provenance record
                    parsed_prov_values = self._parse_prov_test(
                        reg_name, prov_width_list, prov_entry
                    )
                    # if not activity_sent:
                    record_activity(
                        parsed_prov_values["activity"],
                        self.agent_id,
                        parsed_prov_values["version"],
                        "PacketProcessor",
                    )
                    # activity_sent = True
                    # Convert the value if annotation is present
                    converted_value = parsed_prov_values["value"]
                    # if value_annotation:
                    #     converted_value = self._convert_value(
                    #         value_annotation, parsed_prov_values["value"]
                    #     )
                    # Send the provenance event to the message queue
                    record_API_call(
                        reg_name,
                        self.agent_id,
                        parsed_prov_values["version"],
                        "register",
                        parsed_prov_values["operation"],
                        {
                            "index": parsed_prov_values["index"],
                            "value": converted_value,
                            "activity_name": parsed_prov_values["activity"],
                        },
                    )
                    logger.info(
                        f"Register ({reg_name}) entity value: {parsed_prov_values}"
                    )
                else:
                    break
            self.last_count_index[reg_name] += count
            print(
                f"[INFO] Pulling Register ({reg_name}) provenance.. {round(self.last_count_index[reg_name] / 65535, 3) * 100}"
            )
            logger.info(
                f"[INFO] Provenance register capacity: {self.last_count_index[reg_name] / 65535}"
            )
    def _read_packet_in_prov(self):
        # TODO - Modify such that pkt_in details are taken from the details json file.
        # prov_width_list = [
        #     ["ipid", 16],
        #     ["src_ip", 32],
        #     ["activity", 7],
        #     ["version", 8],
        # ]

        # prov_width_list = [
        #     ["src_mac", 16],
        #     ["packet_id", 16],
        #     ["ingress_port", 8],
        #     ["ether_type", 16],
        #     ["version", 3],
        # ]
        prov_width_list = [
            ["packet_id", 16],
            ["version", 8],
        ]

        parsed_prov_values = {}
        total_bit_width = sum([x[1] for x in prov_width_list])

        packet_in_reg = self.controller.read_register("MyIngress.prov_packet_in")
        # packet_in_reg = self.controller.read_register(
        #     "ingress.table0_control.prov_packet_in"
        # )
        # packet_in_reg = self.controller.read_register("prov_packet_in")
        prov_reg_indices = self.controller.read_register("MyIngress.prov_index_reg")
        # prov_reg_indices = self.controller.read_register(
        #     "ingress.table0_control.prov_index_reg"
        # )
        # prov_reg_indices = self.controller.read_register("prov_index_reg")
        end_idx = prov_reg_indices[index_map["packet_in"]]
        start_idx = self.last_count_index["packet_in"]
        idx = start_idx
        entries_read = 0
        # for prov_entry in packet_in_reg[self.last_count_index["packet_in"] :]:
        while idx != end_idx:
            # print("prov entry: ", prov_entry)
            prov_entry = packet_in_reg[idx]
            if prov_entry == 0:
                print(f"Entry 0 - PacketIn: idx: {idx} and end_idx:{end_idx}")
                break
            start_index = 0
            prov_string = format(prov_entry, "b").zfill(total_bit_width)
            # print(f"packetin bit string: ", prov_string)
            for entity_field in prov_width_list:
                entity_type = entity_field[0]
                entity_size = entity_field[1]
                end_index = start_index + entity_size
                entity_value = prov_string[start_index:end_index]
                if entity_type in PKT_IN_TYPE:
                    if PKT_IN_TYPE[entity_type] == "hex":
                        parsed_prov_values[entity_type] = hex(int(entity_value, 2))
                    elif PKT_IN_TYPE[entity_type] == "ip_addr":
                        parsed_prov_values[entity_type] = ".".join(
                            [
                                str((int(entity_value, 2) >> (8 * i)) & 255)
                                for i in range(3, -1, -1)
                            ]
                        )
                    elif PKT_IN_TYPE[entity_type] == "mac_addr":
                        parsed_prov_values[entity_type] = "{:04x}".format(
                            int(entity_value, 2)
                        )
                        parsed_prov_values[entity_type] = ":".join(
                            [
                                parsed_prov_values[entity_type][i : i + 2]
                                for i in range(0, 12, 2)
                            ]
                        )
                    else:
                        logger.warning("PKT_IN_TYPE not recognized")
                else:
                    # Default value type is INT
                    parsed_prov_values[entity_type] = int(entity_value, 2)
                start_index = end_index
                entries_read += 1
                idx = (idx + 1) % REGISTER_SIZE
            logger.info(f"PacketIn values: {parsed_prov_values}")
            self._prepare_prov_api(
                "packet_in", parsed_prov_values
            )  # Send the prov events to the queue
        self.last_count_index["packet_in"] = end_idx
        print(
            f"[INFO] Pulling PacketIN provenance.. {round(self.last_count_index['packet_in'] / REGISTER_SIZE, 3) * 100}"
        )
        logger.info(
            f"[INFO] Provenance register capacity: {self.last_count_index['packet_in'] / 65535}"
        )

    def _read_packet_out_prov(self):
        # TODO - Merge the packet_in and packet_out functions and use a flag

        prov_width_list = [
            ["src_mac", 16],
            ["packet_id", 16],
            ["egress_port", 8],
            ["ether_type", 16],
            ["version", 3],
        ]

        parsed_prov_values = {}
        total_bit_width = sum([x[1] for x in prov_width_list])

        packet_out_reg = self.controller.read_register(
            "ingress.firewall_control.prov_packet_out"
        )
        prov_reg_indices = self.controller.read_register(
            "ingress.firewall_control.prov_index_reg"
        )
        end_idx = prov_reg_indices[index_map["packet_out"]]
        start_idx = self.last_count_index["packet_out"]
        idx = start_idx
        entries_read = 0

        # for prov_entry in packet_in_reg[self.last_count_index["packet_out"] :]:
        while idx != end_idx:
            # print("prov entry: ", prov_entry)
            prov_entry = packet_out_reg[idx]
            if prov_entry == 0:
                print(f"Entry 0 - PacketIn: idx: {idx} and end_idx:{end_idx}")
                break
            # if prov_entry != 0:
            start_index = 0
            prov_string = format(prov_entry, "b").zfill(total_bit_width)
            # print(f"packetin bit string: ", prov_string)
            for entity_field in prov_width_list:
                entity_type = entity_field[0]
                entity_size = entity_field[1]
                end_index = start_index + entity_size
                entity_value = prov_string[start_index:end_index]
                if entity_type in PKT_IN_TYPE:
                    if PKT_IN_TYPE[entity_type] == "hex":
                        parsed_prov_values[entity_type] = hex(int(entity_value, 2))
                    elif PKT_IN_TYPE[entity_type] == "ip_addr":
                        parsed_prov_values[entity_type] = ".".join(
                            [
                                str((int(entity_value, 2) >> (8 * i)) & 255)
                                for i in range(3, -1, -1)
                            ]
                        )
                    elif PKT_IN_TYPE[entity_type] == "mac_addr":
                        parsed_prov_values[entity_type] = "{:04x}".format(
                            int(entity_value, 2)
                        )
                        parsed_prov_values[entity_type] = ":".join(
                            [
                                parsed_prov_values[entity_type][i : i + 2]
                                for i in range(0, 12, 2)
                            ]
                        )
                    else:
                        logger.warning("PktOut TYPE not recognized")
                else:
                    # Default value type is INT
                    parsed_prov_values[entity_type] = int(entity_value, 2)
                start_index = end_index
                entries_read += 1
                idx = (idx + 1) % REGISTER_SIZE
            logger.info(f"PacketOut values: {parsed_prov_values}")
            self._prepare_prov_api(
                "packet_out", parsed_prov_values
            )  # Send the prov events to the queue
            # else:
            self.last_count_index["packet_out"] = end_idx
            print(
                f"[INFO] Pulling PacketOut provenance.. {round(self.last_count_index['packet_out'] / 65535, 3) * 100}"
            )
            # logger.info(
            #     f"[INFO] Provenance register capacity: {self.last_count_index['packet_out'] / 65535}"
            # )

    def _prepare_prov_api(self, entity_name: str, entity_call_values: dict):
        activity_type = "PacketProcessor"
        entity_operation = ""
        version = entity_call_values.pop("version")

        if entity_name == "packet_in":
            # activity_id_str = str(entity_call_values["activity"])
            # activity_name = self.reg_prov_details["packet_in"][activity_id_str]
            activity_name = "send_to_cpu"
            # entity_call_values.pop("activity")
            # entity_call_values.update({"activity_name": activity_name})
            entity_operation = "WRITE"
        elif entity_name == "packet_out":
            activity_name = "set_egress_port"
            entity_call_values.update({"activity_name": activity_name})
            entity_operation = "READ"
        elif entity_name == "table_rule":
            logger.info("table rule API")
            activity_name = "set_egress_port"
        else:
            logger.error("Entity name not recognized...")
        # Use the ProvP4 API calls
        record_agent(self.switch_name, self.agent_id, "Switch", 0)
        record_activity(activity_name, self.agent_id, version, activity_type)
        record_API_call(
            entity_name,
            self.agent_id,
            version,
            entity_name,
            entity_operation,
            {"value": entity_call_values},
        )

    def start(self) -> None:
        """Start the collection of prov events. First prov registers then prov packetIn"""
        POLL_TIME = 2  # In seconds
        try:
            while True:
                self._prov_wraparound()
                # self._read_prov_regs()
                # self._read_packet_in_prov()
                # self._read_packet_out_prov()
                time.sleep(POLL_TIME)
        except KeyboardInterrupt:
            print("Closing DP collector")
            close_channel()


def main():
    """This function 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 as input."
    )
    parser.add_argument("prov_details", type=str, help="Provenance details JSON")
    parser.add_argument("switch_name", type=str, help="Switch Name")

    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

    args = parser.parse_args()
    with open(args.prov_details, "r") as file:
        reg_prov_details = json.load(file)
    collector = ProvCollector(reg_prov_details, thrift_port, args.switch_name)
    collector.start()


if __name__ == "__main__":
    main()
