# This file contains the code to analyze and instrument the P4 program for provenance capture.
import json
import logging
from typing import Any, Set
from collections import defaultdict

from p4_action import P4Action
from p4_register_operation import P4RegisterOperation
# from provp4_augmenter import ProvP4Augmenter

fh_log = logging.FileHandler("provp4_analysis.log")
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s", "%Y-%m-%d")
fh_log.setFormatter(formatter)
logger = logging.getLogger("provp4-analysis")
logger.setLevel(logging.DEBUG)
logger.addHandler(fh_log)

class ProvP4ProgramAnalyzer:
    def __init__(self, compiled_p4: str, prog_name: str):
        self.compiled_p4 = compiled_p4
        self.prog_name = prog_name
        self.reg_bitwidth = {}
        self.fields_size_map = {}
        self.register_list = defaultdict[Any, dict](dict)
        self.augmentation_map = defaultdict[Any, dict](dict)
        self.actions_list = []
        self.json_prov_repr = {}

    def analyze_program(self):
        """
        Main runner function for the P4 program analysis.
        """
        self.extract_size_info()
        self.read_actions, self.write_actions = self.extract_stateful_info()
        self.extract_stateful_actions_info()
        self.write_prov_info_to_file()

    def extract_size_info(self):
        """
        Extract the size information from the P4 program.
        """
        # 1) Register bitwidths
        for register_defs in self.compiled_p4.get("register_arrays", []):
            self.reg_bitwidth[register_defs["name"]] = register_defs["bitwidth"]

        # 2) Header field sizes
        for header in self.compiled_p4.get("header_types", []):
            self.fields_size_map[header["name"]] = {}
            for field in header.get("fields", []):
                self.fields_size_map[header["name"]][field[0]] = field[1]
    def _parse_register_params(self, params, operation):
        if operation == "READ":
            value, name, index = params
        elif operation == "WRITE":
            name, index, value = params
        return name, index, value

    def extract_stateful_info(self):
        """
        Extract the stateful information from the P4 program.
        """
        total_actions, total_reg_writes, total_reg_reads = 0, 0, 0
        write_actions_set = set()
        read_actions_set = set()

        for action in self.compiled_p4["actions"]:
            total_actions += 1
            action_name = action.get("name")
            for primitive in action["primitives"]:
                if primitive["op"] == "register_write":
                    total_reg_writes += 1
                    write_actions_set.add(action_name)
                elif primitive["op"] == "register_read":
                    total_reg_reads += 1
                    read_actions_set.add(action_name)
                # Add support for clone operations in the future
                # elif primitive["op"] == "clone_ingress_pkt_to_egress":
                    # self.clone_details_map[primitive["source_info"]["filename"]].update(
                    #     {
                    #         primitive["source_info"]["line"]: [
                    #             action["id"],
                    #             action["name"],
                    #         ]
                    #     }
                    # )
        read_actions = sorted(read_actions_set)
        write_actions = sorted(write_actions_set)
        logger.info(f"Total Actions : {total_actions}")
        logger.info(f"Total Register writes : {total_reg_writes}")
        logger.info(f"Total Register reads : {total_reg_reads}")
        logger.info(f"Total actions w/ register writes: {len(write_actions)}")
        logger.info(f"Read actions: {read_actions}")
        logger.info(f"Write actions: {write_actions}")
        return read_actions, write_actions

    def extract_stateful_actions_info(self):
        """
        Extract the stateful actions information from the P4 program.
        """
        read_set: Set[str] = set(self.read_actions)
        write_set: Set[str] = set(self.write_actions)
        for action in self.compiled_p4.get("actions", []):
            action_name = action.get("name")
            action_id = action.get("id")
            if not action_name or action_id is None:
            # skip malformed entries
                continue
            needs_reads = action_name in read_set
            needs_writes = action_name in write_set

            current_action = P4Action(action_name, action_id)
            primitives = action.get("primitives", [])
            for prim_idx, primitive in enumerate(primitives):
                op = primitive.get("op")

                # Only process register_read and register_write operations
                if op != "register_read" and op != "register_write":
                    continue
                
                if op == "register_write" and not needs_writes:
                    continue
                if op == "register_read" and not needs_reads:
                    continue
                reg_op = self._build_register_operation(
                    action=action,
                    primitive=primitive,
                    prim_idx=prim_idx,
                    mode="WRITE" if op == "register_write" else "READ",
                )
                if reg_op is None:
                    continue
                # Attach to the correct list on the action object
                if reg_op.op_type == "WRITE":  # adjust attribute name if different
                    current_action.register_writes.append(reg_op)
                else:
                    # Ensure your P4Action has register_reads; if not, create it or use register_ops
                    current_action.register_reads.append(reg_op)
            # Only append if we actually added any operations
            if getattr(current_action, "register_writes", []) or getattr(current_action, "register_reads", []):
                self.actions_list.append(current_action)

    def _build_register_operation(
        self,
        *,
        action,
        primitive,
        prim_idx: int,
        mode: str,  # "READ" or "WRITE"
    ):
        action_name = action.get("name")
        action_id = action.get("id")
        src = primitive.get("source_info", {})

        line = src.get("line")
        filename = src.get("filename")
        fragment = src.get("source_fragment")
        if action_name is None or action_id is None or line is None:
            return None
        uid = int(f"{action_id}{line}") # Use a hash funtion
        name, index, value = self._parse_register_params(primitive.get("parameters", []), mode)
        reg_name = name.get("value")
        reg_name_key = reg_name.replace(".", "_") # P4 does not allow . in register names
        if not reg_name:
            return None
        # Keep track of which actions touch which registers
        self.register_list[reg_name_key].update({action_id: action_name})
        bitwidth = self.reg_bitwidth.get(reg_name)
        if bitwidth is None:
            raise Exception("Bitwidth not found in compiled JSON")
        current_reg_op = P4RegisterOperation(
            mode,
            reg_name_key,
            action_name,
            action_id,
            uid,
            index,
            value,
            fragment,
            line,
            bitwidth,
            filename,
        )
        # Compute sizes and update register_list
        current_reg_op.compute_prov_size(self.fields_size_map)
        if current_reg_op.reg_size_exceeded:
            # Recompute the sizes
            current_reg_op.compute_prov_size(self.fields_size_map)
        if current_reg_op.index_size is None or current_reg_op.value_size is None:
            raise Exception("Index or value size not found in compiled JSON")
        if current_reg_op.index_size == 0 or current_reg_op.value_size == 0:
            raise Exception("Index or value size is 0")
        self.register_list[reg_name_key].update(
            {"index_size": current_reg_op.index_size, "value_size": current_reg_op.value_size}
        )

        # Provenance definitions + JSON representation for augmentation
        current_reg_op.make_prov_definition()
        current_reg_op.prov_register_declaration()
        self.json_prov_repr.update(current_reg_op.get_json_repr())
        print(self.json_prov_repr)
        # Save augmentation info
        self._save_for_augmentation(
            filename,
            line,
            current_reg_op.uid,
            current_reg_op.log_write_str,
            current_reg_op.reg_decl,
            current_reg_op.prov_index_str,
        )
        # print(self.augmentation_map)
        return current_reg_op

    def _save_for_augmentation(
        self,
        filename: str,
        line_no: str,
        statement_id: int,
        write_str: str,
        reg_decl_str: str,
        prov_index_str: str,
    ):
        self.augmentation_map[filename].update(
            {line_no: [statement_id, write_str, reg_decl_str, prov_index_str]}
        )
    def write_prov_info_to_file(self):
        """
        Write the provenance information to a file.
        """
        # TODO - Add PacketIn and PacketOut information later
        with open(f"{self.prog_name}_prov_info.json", "w") as file:
            json.dump(self.register_list, file)

# if __name__ == "__main__":
    