from __future__ import annotations

from typing import TYPE_CHECKING

from archinfo.arch_arm import get_real_address_if_arm

from angrmanagement.data.object_container import ObjectContainer

if TYPE_CHECKING:
    from angr.sim_variable import SimVariable


class OperandHighlightMode:
    SAME_IDENT = 0
    SAME_TEXT = 1


class OperandDescriptor:
    __slots__ = (
        "text",
        "num_value",
        "func_addr",
        "variable_ident",
    )

    def __init__(self, text: str, num_value, func_addr=None, variable_ident=None) -> None:
        self.text = text
        self.num_value = num_value
        self.func_addr = func_addr
        self.variable_ident = variable_ident


class InfoDock:
    """
    Stores information associated to a disassembly view. Such information will be shared between the graph view and the
    linear view.
    """

    def __init__(self, disasm_view) -> None:
        super().__init__()
        self.disasm_view = disasm_view

        self.induction_variable_analysis = None
        self.variable_manager = None

        self.highlight_mode = OperandHighlightMode.SAME_IDENT  # default highlight mode

        self.selected_insns = ObjectContainer(set(), "The currently selected instructions")
        self.selected_operands = ObjectContainer({}, "The currently selected instruction operands")
        self.selected_blocks = ObjectContainer(set(), "The currently selected blocks")
        self.selected_ail_statements = ObjectContainer(set(), "The currently selected AIL statements")
        self.hovered_block = ObjectContainer(None, "The currently hovered block")
        self.hovered_edge = ObjectContainer(None, "The currently hovered edge")
        self.selected_labels = ObjectContainer(set(), "The currently selected labels")
        self.selected_variables = ObjectContainer(set(), "The currently selected variables")
        self.selected_block_tree_node = ObjectContainer(None, "The currently selected BlockTreeNode object")

    @property
    def smart_highlighting(self):
        return self.highlight_mode == OperandHighlightMode.SAME_IDENT

    @smart_highlighting.setter
    def smart_highlighting(self, v) -> None:
        if v:
            self.highlight_mode = OperandHighlightMode.SAME_IDENT
        else:
            self.highlight_mode = OperandHighlightMode.SAME_TEXT

    def initialize(self) -> None:
        self.selected_blocks.clear()
        self.selected_insns.clear()
        self.selected_operands.clear()
        self.hovered_block.am_obj = None

    def copy(self) -> InfoDock:
        r = InfoDock(self.disasm_view)
        r.variable_manager = self.variable_manager
        r.highlight_mode = self.highlight_mode
        r.selected_insns.am_obj = set(self.selected_insns.am_obj)
        r.selected_operands.am_obj = dict(self.selected_operands.am_obj)
        r.selected_blocks.am_obj = set(self.selected_blocks.am_obj)
        r.hovered_block.am_obj = self.hovered_block.am_obj
        r.hovered_edge.am_obj = self.hovered_edge.am_obj
        r.selected_labels.am_obj = set(self.selected_labels.am_obj)
        return r

    def hover_edge(self, src_addr, dst_addr) -> None:
        self.hovered_edge.am_obj = src_addr, dst_addr
        self.hovered_edge.am_event()

    def unhover_edge(self, src_addr, dst_addr) -> None:
        if self.hovered_edge.am_obj == (src_addr, dst_addr):
            self.hovered_edge.am_obj = None
            self.hovered_edge.am_event()

    def hover_block(self, block_addr) -> None:
        self.hovered_block.am_obj = block_addr
        self.hovered_block.am_event()

    def unhover_block(self, block_addr) -> None:
        if self.hovered_block.am_obj == block_addr:
            self.hovered_block.am_obj = None
            self.hovered_block.am_event()

    def select_block(self, block_addr) -> None:
        self.selected_blocks.clear()  # selecting one block at a time
        self.selected_blocks.add(block_addr)
        self.selected_blocks.am_event()
        self._update_published_view_state()

    def select_instruction(self, insn_addr, unique: bool = True, insn_pos=None, use_animation: bool = True) -> None:
        arch = self.disasm_view.instance.project.arch
        real_addr = get_real_address_if_arm(arch, insn_addr)
        self.disasm_view.set_synchronized_cursor_address(real_addr)

        self.unselect_all_labels()
        self.unselect_block_tree_node()
        if insn_addr not in self.selected_insns:
            if unique:
                # unselect existing ones
                self.unselect_all_instructions()
                self.selected_insns.add(insn_addr)
            else:
                self.selected_insns.add(insn_addr)
            self.disasm_view.current_graph.show_instruction(insn_addr, insn_pos=insn_pos, use_animation=use_animation)
            self.selected_insns.am_event(insn_addr=insn_addr)

        self._update_published_view_state()

    def unselect_instruction(self, insn_addr) -> None:
        if insn_addr in self.selected_insns:
            self.selected_insns.remove(insn_addr)
            self.selected_insns.am_event()
        self._update_published_view_state()

    def unselect_all_instructions(self) -> None:
        if self.selected_insns:
            self.selected_insns.clear()
            self.selected_insns.am_event()
        self._update_published_view_state()

    def select_ail_statement(self, stmt_addr, unique: bool = True) -> None:
        self.disasm_view.set_synchronized_cursor_address(stmt_addr[0])

        self.unselect_all_labels()
        self.unselect_all_instructions()
        if stmt_addr not in self.selected_ail_statements:
            if unique:
                # unselect existing ones
                self.unselect_all_instructions()
                self.selected_ail_statements.add(stmt_addr)
            else:
                self.selected_ail_statements.add(stmt_addr)
            self.selected_ail_statements.am_event(stmt_addr=stmt_addr)

        self._update_published_view_state()

    def unselect_ail_statement(self, stmt_addr) -> None:
        if stmt_addr in self.selected_ail_statements:
            self.selected_ail_statements.remove(stmt_addr)
            self.selected_ail_statements.am_event()
        self._update_published_view_state()

    def unselect_all_ail_statements(self) -> None:
        if self.selected_ail_statements:
            self.selected_ail_statements.clear()
            self.selected_ail_statements.am_event()
        self._update_published_view_state()

    def select_operand(
        self, ins_addr: int, operand_index: int, operand: OperandDescriptor, unique: bool = False
    ) -> None:
        """
        Mark an operand as selected.

        :param ins_addr:                Address of the instruction.
        :param operand_index:           Index of the operand.
        :param operand:   Data of the operand.
        :param unique:                 If this is a unique selection or not.
        :return:                            None
        """

        tpl = ins_addr, operand_index
        if tpl not in self.selected_operands:
            if unique:
                self.selected_operands.clear()
            self.selected_operands[tpl] = operand
            self.selected_operands.am_event()

    def unselect_operand(self, insn_addr, operand_idx) -> None:
        if (insn_addr, operand_idx) in self.selected_operands:
            self.selected_operands.pop((insn_addr, operand_idx))
            self.selected_operands.am_event()

    def unselect_all_operands(self) -> None:
        if self.selected_operands:
            self.selected_operands.clear()
            self.selected_operands.am_event()

    def select_label(self, label_addr) -> None:
        arch = self.disasm_view.instance.project.arch
        real_addr = get_real_address_if_arm(arch, label_addr)
        self.disasm_view.set_synchronized_cursor_address(real_addr)

        # only one label can be selected at a time
        # also, clear selection of instructions and operands
        self.unselect_all_instructions()
        self.unselect_all_operands()
        self.unselect_block_tree_node()

        self.selected_labels.clear()
        self.selected_labels.add(label_addr)
        self.selected_labels.am_event()

        self._update_published_view_state()

    def toggle_label_selection(self, addr: int) -> None:
        """
        Toggle the selection state of a label in the disassembly view.

        :param addr:    Address of the instruction to toggle.
        """

        if addr in self.selected_labels:
            self.unselect_label(addr)
        else:
            self.select_label(addr)

    def unselect_label(self, label_addr) -> None:
        if label_addr in self.selected_labels:
            self.selected_labels.remove(label_addr)
            self.selected_labels.am_event()
        self._update_published_view_state()

    def unselect_all_labels(self) -> None:
        self.selected_labels.clear()
        self.selected_labels.am_event()
        self._update_published_view_state()

    def toggle_instruction_selection(self, insn_addr, insn_pos=None, unique: bool = False) -> None:
        """
        Toggle the selection state of an instruction in the disassembly view.

        :param int insn_addr: Address of the instruction to toggle.
        :return:              None
        """

        if insn_addr in self.selected_insns:
            self.unselect_instruction(insn_addr)
        else:
            self.select_instruction(insn_addr, unique=unique, insn_pos=insn_pos)

    def toggle_operand_selection(self, insn_addr, operand_idx, operand, insn_pos=None, unique: bool = False) -> bool:
        """
        Toggle the selection state of an operand of an instruction in the disassembly view.

        :param int insn_addr:   Address of the instruction to toggle.
        :param int operand_idx: The operand to toggle.
        :param operand:         The operand instance.
        :return:                True if this operand is now selected, False otherwise.
        :rtype:                 bool
        """

        if (insn_addr, operand_idx) in self.selected_operands:
            self.unselect_operand(insn_addr, operand_idx)
            return False
        else:
            self.select_operand(insn_addr, operand_idx, operand, unique=unique)
            self.disasm_view.current_graph.show_instruction(insn_addr, insn_pos=insn_pos)
            return True

    def select_variable(self, unified_variable: SimVariable, unique: bool = True) -> None:
        self.unselect_all_labels()
        self.unselect_all_instructions()
        self.unselect_all_operands()

        if unique:
            self.selected_variables.clear()
        if unified_variable not in self.selected_variables:
            self.selected_variables.add(unified_variable)
            self.selected_variables.am_event()

    def unselect_variable(self, unified_variable: SimVariable) -> None:
        if unified_variable in self.selected_variables:
            self.selected_variables.remove(unified_variable)
            self.selected_variables.am_event()

    def toggle_variable_selection(self, unified_variable: SimVariable, unique: bool = True) -> bool:
        if len(self.selected_variables) > 1 and unique:
            # multiple variables are selected
            # clear existing selections and select this one
            self.select_variable(unified_variable, unique=True)
            return True

        if unified_variable in self.selected_variables:
            self.unselect_variable(unified_variable)
            return False
        else:
            self.select_variable(unified_variable, unique=unique)
            return True

    def clear_selection(self) -> None:
        self.selected_blocks.clear()
        self.selected_blocks.am_event()

        self.selected_insns.clear()
        self.selected_insns.am_event()

        self.selected_operands.clear()
        self.selected_operands.am_event()

        self.selected_variables.clear()
        self.selected_variables.am_event()

        self.selected_block_tree_node.am_obj = None
        self.selected_block_tree_node.am_event()

        self._update_published_view_state()

    def is_edge_hovered(self, src_addr, dst_addr):
        return self.hovered_edge.am_obj == (src_addr, dst_addr)

    def is_block_hovered(self, block_addr):
        return block_addr == self.hovered_block.am_obj

    def is_block_selected(self, block_addr) -> bool:
        return block_addr in self.selected_blocks

    def is_instruction_selected(self, ins_addr) -> bool:
        """
        Check if an instruction at @ins_addr is currently selected or not.

        :param int ins_addr:    Address of the instruction.
        :return:                True if it is selected, False otherwise.
        :rtype:                 bool
        """
        return ins_addr in self.selected_insns

    def is_operand_selected(self, ins_addr, operand_index) -> bool:
        """
        Check if an operand at @ins_addr and @operand_index is currently selected or not.

        :param int ins_addr:        Address of the instruction
        :param int operand_index:   Index of the operand.
        :return:                    bool
        """
        return (ins_addr, operand_index) in self.selected_operands

    def is_label_selected(self, label_addr) -> bool:
        return label_addr in self.selected_labels

    def is_variable_selected(self, unique_variable_ident: str) -> bool:
        return unique_variable_ident in self.selected_variables

    def select_block_tree_node(self, obj) -> None:
        """
        For QBlockCodeObj, we simply track selected state for now and handle
        matching in object handlers
        """

        self.unselect_all_instructions()
        self.unselect_all_labels()

        self.selected_block_tree_node.am_obj = obj
        self.selected_block_tree_node.am_event()

    def unselect_block_tree_node(self):
        self.selected_block_tree_node.am_obj = None
        self.selected_block_tree_node.am_event()

    def toggle_block_tree_node_selection(self, obj) -> None:
        """
        Toggle the selection state of a QBlockCodeObj in the disassembly view.

        :param obj:    The QBlockCodeObj to toggle.
        """

        if self.selected_block_tree_node.am_obj == obj:
            self.unselect_block_tree_node()
        else:
            self.select_block_tree_node(obj)

    def _update_published_view_state(self) -> None:
        self.disasm_view.published_view_state.cursors = self.selected_insns.union(self.selected_labels)
        self.disasm_view.notify_view_state_updated()
