ACTIVITY_SIZE = 8 # 8 bits for activity ID
OPERATION_SIZE = 1 # 1 bit for READ/WRITE tag
EPOCH_SIZE = 8 # 8 bits for epoch
AVAILABLE_PROV_WIDTH = 64 - ACTIVITY_SIZE - OPERATION_SIZE
class P4RegisterOperation:
    def __init__(
        self,
        op_type: str,
        reg_name: str,
        action_name: str,
        action_id: str,
        uid: int,
        index: dict,
        value: str,
        src_fragment: str,
        line_no: int,
        bitwidth: int,
        src_filename: str,
    ) -> None:
        self.op_type = op_type
        self.reg_name = reg_name.replace(".", "_") # P4 does not allow . in register names
        self.action_name = action_name
        self.action_id = action_id
        self.uid = uid
        self.index = index
        self.value = value
        self.src_fragment = src_fragment
        self.line_no = line_no
        self.index_size = 0
        self.value_size = bitwidth
        self.src_file_name = src_filename
        self.log_write_str = ""
        # self.PROV_REG_DEF = (
        #     f"read_{self.uid}" if self.op_type == "READ" else f"write_{self.uid}"
        # )
        self.PROV_REG_DEF = f"prov_{self.reg_name}"
        self.PROV_INDEX_WIDTH = 32
        self.prov_width = 0
        self.reg_decl = ""
        self.json_repr = {}
        self.total_prov_width = 0
        self.reg_size_exceeded = False
        self.available_prov_width = AVAILABLE_PROV_WIDTH
    # TODO - Fix the provenance register names for augmentation for SECAP

    def __str__(self) -> str:
        return f"Reg Name - {self.reg_name}, UID: {self.uid}"

    def compute_prov_size(self, fields_size_map: dict) -> None:
        if self.reg_size_exceeded:
            return
        if (self.index["type"] == "field") and ("scalars" in self.index["value"]):
            self.index_size = fields_size_map["scalars_0"][self.index["value"][1]]
            self.prov_width = self.index_size + self.value_size
            self.total_prov_width = (
                self.prov_width + ACTIVITY_SIZE + OPERATION_SIZE + EPOCH_SIZE
            )
        elif self.index["type"] == "hexstr":
            hex_str = self.index["value"]
            if hex_str.startswith("0x"):
                hex_digits = hex_str[2:].lstrip("0") or "0"
            else:
                hex_digits = hex_str.lstrip("0") or "0"
            self.index_size = len(hex_digits) * 4
            self.prov_width = self.index_size + self.value_size
            self.total_prov_width = (
                self.prov_width + ACTIVITY_SIZE + OPERATION_SIZE + EPOCH_SIZE
            )
        elif (
            self.index["type"] == "field"
            and self.index["value"][0] == "standard_metadata"
        ):
            self.index_size = fields_size_map["standard_metadata"][
                self.index["value"][1]
            ]
            self.prov_width = self.index_size + self.value_size
            self.total_prov_width = self.prov_width + ACTIVITY_SIZE + OPERATION_SIZE + EPOCH_SIZE
        else:
            print(self.index)
            raise NotImplementedError(
                f"{self.index['value']}-Index not in scalars. Check index value"
            )

    def make_prov_definition(self) -> None:
        # TODO - Metadata variable name needs to be read from the program - Programmer defined
        PROV_INDEX_STR = f"(bit<32>) meta.prov_index_val "
        PROV_INDEX_NAME = f"INDEX_{self.reg_name}"
        log_string = f"\t\tget_prov_index({PROV_INDEX_NAME});\n"
        action_tag = f"(bit<8>) {self.action_id}"
        if self.op_type == "WRITE":
            src_write_args = self.src_fragment.split(".write")[1]
            operation_tag = "(bit<1>) 1"
        elif self.op_type == "READ":
            src_write_args = self.src_fragment.split(".read")[1]
            operation_tag = "(bit<1>) 0"
        epoch_tag = "(bit<8>)meta.current_epoch"
        log_string += f"\t\t{self.PROV_REG_DEF}.write({PROV_INDEX_STR}, {src_write_args[1:-1].replace(',', ' ++')} ++ {action_tag} ++ {epoch_tag} ++ {operation_tag});"
        log_string += f"\n\t\tincrement_prov_index({PROV_INDEX_NAME});\n"
        self.log_write_str = log_string
        self.prov_index_str = PROV_INDEX_NAME

    def prov_register_declaration(self) -> None:
        # TODO - Change the 65535 to 2^(width)
        if self.total_prov_width > AVAILABLE_PROV_WIDTH:
            self.reg_size_exceeded = True
            print(f"[Warning] Provenance register size for register: {self.reg_name} ({self.total_prov_width}) exceeded {AVAILABLE_PROV_WIDTH}...will be truncated to {AVAILABLE_PROV_WIDTH} bits")
            value_size_original = self.value_size
            self.value_size = AVAILABLE_PROV_WIDTH - self.index_size
            self.total_prov_width = AVAILABLE_PROV_WIDTH
            print(f"Truncated value size for register: {self.reg_name} to {self.value_size} bits from {value_size_original} bits to {AVAILABLE_PROV_WIDTH} bits")
        self.reg_decl = f"register<bit<{self.total_prov_width}>, bit<{self.PROV_INDEX_WIDTH}>>(65535) {self.PROV_REG_DEF};\n"

    def get_json_repr(self) -> dict:
        self.json_repr[self.PROV_REG_DEF] = [
            self.reg_name,
            self.action_name,
            self.index_size,
            self.value_size,
        ]
        return self.json_repr
