from __future__ import annotations
from typing import Literal, TYPE_CHECKING, overload
import logging
from collections import defaultdict
from itertools import count, chain

from sortedcontainers import SortedDict
import networkx

import angr.ailment as ailment
from cle.backends.elf.compilation_unit import CompilationUnit
from cle.backends.elf.variable import Variable

from angr.utils.orderedset import OrderedSet
from angr.utils.ail import is_phi_assignment
from angr.utils.types import unpack_pointer, replace_pointer_pts_to
from angr.protos import variables_pb2
from angr.serializable import Serializable
from angr.sim_variable import SimVariable, SimStackVariable, SimMemoryVariable, SimRegisterVariable
from angr.sim_type import (
    TypeRef,
    SimType,
    SimStruct,
    SimTypeBottom,
    SimTypeChar,
    SimTypeShort,
    SimTypeInt,
    SimTypeLong,
)
from angr.keyed_region import KeyedRegion
from angr.knowledge_plugins.plugin import KnowledgeBasePlugin
from angr.knowledge_plugins.types import TypesStore
from .variable_access import VariableAccess, VariableAccessSort

if TYPE_CHECKING:
    from angr.analyses.decompiler.stack_item import StackItem
    from angr.code_location import CodeLocation

l = logging.getLogger(name=__name__)


class VariableType:
    """
    Describes variable types.
    """

    REGISTER = 0
    MEMORY = 1


class LiveVariables:
    """
    A collection of live variables at a program point.
    """

    __slots__ = (
        "register_region",
        "stack_region",
    )

    def __init__(self, register_region, stack_region):
        self.register_region = register_region
        self.stack_region = stack_region


def _defaultdict_set():
    return defaultdict(set)


class VariableManagerInternal(Serializable):
    """
    Manage variables for a function. It is meant to be used internally by VariableManager, but it's common to be
    given a reference to one in response to a query for "the variables for a given function". Maybe a better name
    would be "VariableManagerScope".
    """

    def __init__(self, manager, func_addr=None):
        self.manager: VariableManager = manager

        self.func_addr = func_addr

        self._variables: OrderedSet[SimVariable] = OrderedSet()  # all variables that are added to any region
        self._global_region = KeyedRegion()
        self._stack_region = KeyedRegion()
        self._register_region = KeyedRegion()
        self._live_variables = {}  # a mapping between addresses of program points and live variable collections

        self._variable_accesses: dict[SimVariable, set[VariableAccess]] = defaultdict(set)
        self._insn_to_variable: dict[int, set[tuple[SimVariable, int | None]]] = defaultdict(set)
        self._stmt_to_variable: dict[tuple[int, int] | tuple[int, int, int], set[tuple[SimVariable, int]]] = (
            defaultdict(set)
        )
        self._variable_to_stmt: dict[SimVariable, set[tuple[int, int] | tuple[int, int, int]]] = defaultdict(set)
        self._atom_to_variable: dict[
            tuple[int, int] | tuple[int, int, int], dict[int, set[tuple[SimVariable, int]]]
        ] = defaultdict(_defaultdict_set)
        self._ident_to_variable: dict[str, SimVariable] = {}
        self._vvarid_to_variable: dict[int, SimVariable] = {}
        self._variable_to_vvarids: dict[SimVariable, set[int]] = defaultdict(set)
        self._variable_counters = {
            "register": count(),
            "stack": count(),
            "argument": count(),
            "phi": count(),
            "global": count(),
        }

        self._unified_variables: set[SimVariable] = set()
        self._variables_to_unified_variables: dict[SimVariable, SimVariable] = {}

        self._phi_variables = {}
        self._variables_to_phivars = defaultdict(set)
        self._phi_variables_by_block = defaultdict(set)

        self.types = TypesStore(self.manager._kb)
        self.variable_to_types: dict[SimVariable, SimType] = {}
        self.variables_with_manual_types = set()

        # optimization
        self._variables_without_writes = set()

        # dict[int, tuple[SimStackVariable, SimStruct]]
        self.stack_offset_to_struct = SortedDict()

        self.ret_val_size = None

    #
    # Serialization
    #

    def __setstate__(self, state):
        self.__dict__.update(state)

    def __getstate__(self):
        attributes = [
            "func_addr",
            "_variables",
            "_global_region",
            "_stack_region",
            "_register_region",
            "_live_variables",
            "_variable_accesses",
            "_insn_to_variable",
            "_stmt_to_variable",
            "_variable_to_stmt",
            "_atom_to_variable",
            "_ident_to_variable",
            "_vvarid_to_variable",
            "_variable_to_vvarids",
            "_variable_counters",
            "_unified_variables",
            "_variables_to_unified_variables",
            "_phi_variables",
            "_variables_to_phivars",
            "_phi_variables_by_block",
            "types",
            "variable_to_types",
            "variables_with_manual_types",
            "_variables_without_writes",
            "ret_val_size",
        ]
        d = {k: getattr(self, k) for k in attributes}
        d["manager"] = None
        d["types"]._kb = None
        return d

    def set_manager(self, manager: VariableManager):
        self.manager = manager
        self.types._kb = manager._kb

    @classmethod
    def _get_cmsg(cls):
        return (
            variables_pb2.VariableManagerInternal()  # type: ignore[reportAttributeAccessIssue]  # pylint:disable=no-member
        )

    def serialize_to_cmessage(self):
        # pylint:disable=no-member,unused-variable
        cmsg = self._get_cmsg()

        # variables
        register_variables = []
        stack_variables = []
        memory_variables = []

        for variable in self._variables:
            vc = variable.serialize_to_cmessage()
            if isinstance(variable, SimRegisterVariable):
                register_variables.append(vc)
            elif isinstance(variable, SimStackVariable):
                stack_variables.append(vc)
            elif isinstance(variable, SimMemoryVariable):
                memory_variables.append(vc)
            else:
                raise NotImplementedError
        for variable in self._phi_variables:
            vc = variable.serialize_to_cmessage()
            vc.base.is_phi = True
            if isinstance(variable, SimRegisterVariable):
                register_variables.append(vc)
            elif isinstance(variable, SimStackVariable):
                stack_variables.append(vc)
            elif isinstance(variable, SimMemoryVariable):
                memory_variables.append(vc)
            else:
                raise NotImplementedError

        cmsg.regvars.extend(register_variables)
        cmsg.stackvars.extend(stack_variables)
        cmsg.memvars.extend(memory_variables)

        # accesses
        accesses = []
        for variable_accesses in self._variable_accesses.values():
            for variable_access in variable_accesses:
                accesses.append(variable_access.serialize_to_cmessage())
        cmsg.accesses.extend(accesses)

        # unified variables
        unified_register_variables = []
        unified_stack_variables = []
        unified_memory_variables = []

        unified_variable_idents: set[str] = set()
        for variable in self._unified_variables:
            assert isinstance(variable.ident, str)
            unified_variable_idents.add(variable.ident)
            if isinstance(variable, SimRegisterVariable):
                unified_register_variables.append(variable.serialize_to_cmessage())
            elif isinstance(variable, SimStackVariable):
                unified_stack_variables.append(variable.serialize_to_cmessage())
            elif isinstance(variable, SimMemoryVariable):
                unified_memory_variables.append(variable.serialize_to_cmessage())
            else:
                raise NotImplementedError

        cmsg.unified_regvars.extend(unified_register_variables)
        cmsg.unified_stackvars.extend(unified_stack_variables)
        cmsg.unified_memvars.extend(unified_memory_variables)

        relations = []
        for variable, unified in self._variables_to_unified_variables.items():
            if unified.ident not in unified_variable_idents:
                l.error(
                    "The unified variable %s is missing from the unified variables of function %#x. Please "
                    "report it on GitHub.",
                    unified.ident,
                    self.func_addr,
                )
                continue
            relation = variables_pb2.Var2Unified()  # type: ignore[reportAttributeAccessIssue]
            relation.var_ident = variable.ident
            relation.unified_var_ident = unified.ident
            relations.append(relation)
        cmsg.var2unified.extend(relations)

        # phi vars
        phi_relations = []
        for phi, vars_ in self._phi_variables.items():
            for var in vars_:
                if var not in self._variables and var not in self._phi_variables:
                    l.error("Ignore variable %s because it is not in the registered list.", var.ident)
                    continue
                relation = variables_pb2.Phi2Var()  # type: ignore[reportAttributeAccessIssue]
                relation.phi_ident = phi.ident
                relation.var_ident = var.ident
                phi_relations.append(relation)
        cmsg.phi2var.extend(phi_relations)

        # TODO: Types

        # TODO: vvarid_to_varialbes & variable_to_vvarids

        return cmsg

    @classmethod
    def parse_from_cmessage(
        cls, cmsg, variable_manager=None, func_addr=None, **kwargs
    ) -> VariableManagerInternal:  # pylint:disable=arguments-differ
        model = VariableManagerInternal(variable_manager, func_addr=func_addr)

        variable_by_ident = {}

        # variables
        all_vars = []

        for regvar_pb2 in cmsg.regvars:
            all_vars.append(
                (
                    regvar_pb2.base.is_phi,  # type: ignore[reportAttributeAccessIssue]
                    SimRegisterVariable.parse_from_cmessage(regvar_pb2),
                )
            )
        for stackvar_pb2 in cmsg.stackvars:
            all_vars.append(
                (
                    stackvar_pb2.base.is_phi,  # type: ignore[reportAttributeAccessIssue]
                    SimStackVariable.parse_from_cmessage(stackvar_pb2),
                )
            )
        for memvar_pb2 in cmsg.memvars:
            all_vars.append(
                (
                    memvar_pb2.base.is_phi,  # type: ignore[reportAttributeAccessIssue]
                    SimMemoryVariable.parse_from_cmessage(memvar_pb2),
                )
            )
        for is_phi, var in all_vars:
            variable_by_ident[var.ident] = var
            if is_phi:
                model._phi_variables[var] = set()
            else:
                model._variables.add(var)
                model._ident_to_variable[var.ident] = var

        # variable accesses
        for varaccess_pb2 in cmsg.accesses:
            variable_access = VariableAccess.parse_from_cmessage(varaccess_pb2, variable_by_ident=variable_by_ident)
            variable = variable_access.variable
            offset = variable_access.offset
            assert variable is not None and offset is not None
            tpl = (variable, offset)

            model._variable_accesses[variable_access.variable].add(variable_access)
            assert variable_access.location.ins_addr is not None
            model._insn_to_variable[variable_access.location.ins_addr].add(tpl)
            assert variable_access.location.block_addr is not None
            assert variable_access.location.stmt_idx is not None
            loc = (
                (variable_access.location.block_addr, variable_access.location.stmt_idx)
                if variable_access.location.block_idx is None
                else (
                    variable_access.location.block_addr,
                    variable_access.location.block_idx,
                    variable_access.location.stmt_idx,
                )
            )
            model._stmt_to_variable[loc].add(tpl)
            model._variable_to_stmt[variable].add(loc)
            if variable_access.atom_hash is not None:
                model._atom_to_variable[loc][variable_access.atom_hash].add(tpl)

        # unified variables
        unified_variable_by_ident = {}
        for regvar_pb2 in cmsg.unified_regvars:
            regvar = SimRegisterVariable.parse_from_cmessage(regvar_pb2)
            unified_variable_by_ident[regvar.ident] = regvar
            model._unified_variables.add(regvar)
        for stackvar_pb2 in cmsg.unified_stackvars:
            stackvar = SimStackVariable.parse_from_cmessage(stackvar_pb2)
            unified_variable_by_ident[stackvar.ident] = stackvar
            model._unified_variables.add(stackvar)
        for memvar_pb2 in cmsg.unified_memvars:
            memvar = SimMemoryVariable.parse_from_cmessage(memvar_pb2)
            unified_variable_by_ident[memvar.ident] = memvar
            model._unified_variables.add(memvar)

        for var2unified in cmsg.var2unified:
            variable = variable_by_ident[var2unified.var_ident]
            unified = unified_variable_by_ident.get(var2unified.unified_var_ident, None)
            if unified is None:
                l.warning(
                    "Unified variable %s is not found in unified_variable_by_ident.", var2unified.unified_var_ident
                )
                # as a stop gap, we make the variable unify to itself
                model._variables_to_unified_variables[variable] = variable
                continue
            model._variables_to_unified_variables[variable] = unified

        for phi2var in cmsg.phi2var:
            phi = variable_by_ident.get(phi2var.phi_ident, None)
            if phi is None:
                l.warning("Phi variable %s is not found in variable_by_ident.", phi2var.phi_ident)
                continue
            var = variable_by_ident.get(phi2var.var_ident, None)
            if var is None:
                l.warning("Variable %s is not found in variable_by_ident.", phi2var.var_ident)
                continue
            model._phi_variables[phi].add(var)
            model._variables_to_phivars[var].add(phi)

        # TODO: Types

        for var in model._variables:
            if isinstance(var, SimStackVariable):
                region = model._stack_region
                offset = var.offset
            elif isinstance(var, SimRegisterVariable):
                region = model._register_region
                offset = var.reg
            elif isinstance(var, SimMemoryVariable):
                region = model._global_region
                offset = var.addr
            else:
                raise ValueError(f"Unsupported sort {type(var)} in parse_from_cmessage().")

            region.add_variable(offset, var)

        model._variables_without_writes = set(model.get_variables_without_writes())

        return model

    #
    # Public methods
    #

    def next_variable_ident(self, sort):
        if sort not in self._variable_counters:
            raise ValueError(f"Unsupported variable sort {sort}")

        if sort == "register":
            prefix = "r"
        elif sort == "stack":
            prefix = "s"
        elif sort == "argument":
            prefix = "arg"
        elif sort == "global":
            prefix = "g"
        else:
            prefix = "m"

        return f"i{prefix}_{next(self._variable_counters[sort])}"

    def add_variable(self, sort, start, variable: SimVariable):
        if sort == "stack":
            region = self._stack_region
        elif sort == "register":
            region = self._register_region
        elif sort == "global":
            region = self._global_region
        else:
            raise ValueError(f"Unsupported sort {sort} in add_variable().")

        if variable.ident is not None:
            # find if there is already an existing variable with the same identifier
            if variable.ident in self._ident_to_variable:
                existing_var = self._ident_to_variable[variable.ident]
                if existing_var.name is not None and not variable.renamed:
                    variable.name = existing_var.name
                    variable.renamed = existing_var.renamed
            self._ident_to_variable[variable.ident] = variable

        region.add_variable(start, variable)
        self._variables.add(variable)
        self._variables_without_writes.add(variable)

    def set_variable(self, sort, start, variable: SimVariable):
        if sort == "stack":
            region = self._stack_region
        elif sort == "register":
            region = self._register_region
        elif sort == "global":
            region = self._global_region
        else:
            raise ValueError(f"Unsupported sort {sort} in set_variable().")
        # find if there is already an existing variable with the same identifier
        if variable.ident in self._ident_to_variable:
            existing_var = self._ident_to_variable[variable.ident]
            if existing_var.name is not None and not variable.renamed:
                variable.name = existing_var.name
                variable.renamed = existing_var.renamed
        region.set_variable(start, variable)
        self._variables.add(variable)
        self._variables_without_writes.add(variable)

    def write_to(self, variable, offset, location, overwrite=False, atom=None):
        self._record_variable_access(
            VariableAccessSort.WRITE, variable, offset, location, overwrite=overwrite, atom=atom
        )

    def read_from(self, variable, offset, location, overwrite=False, atom=None):
        self._record_variable_access(
            VariableAccessSort.READ, variable, offset, location, overwrite=overwrite, atom=atom
        )

    def reference_at(self, variable, offset, location, overwrite=False, atom=None):
        self._record_variable_access(
            VariableAccessSort.REFERENCE, variable, offset, location, overwrite=overwrite, atom=atom
        )

    def _record_variable_access(
        self,
        sort: int,
        variable,
        offset,
        location: CodeLocation,
        overwrite=False,
        atom=None,
    ):
        atom_hash = (hash(atom) & 0xFFFF_FFFF) if atom is not None else None
        if overwrite:
            self._variable_accesses[variable] = {VariableAccess(variable, sort, location, offset, atom_hash=atom_hash)}
        else:
            self._variable_accesses[variable].add(VariableAccess(variable, sort, location, offset, atom_hash=atom_hash))
        self.record_variable(location, variable, offset, overwrite=overwrite, atom=atom)
        if sort == VariableAccessSort.WRITE and variable in self._variables_without_writes:
            self._variables_without_writes.discard(variable)

    def record_variable(self, location: CodeLocation, variable, offset, overwrite=False, atom=None):
        if variable.ident not in self._ident_to_variable:
            self._ident_to_variable[variable.ident] = variable
            self._variables.add(variable)
        var_and_offset = variable, offset
        atom_hash = (hash(atom) & 0xFFFF_FFFF) if atom is not None else None
        assert location.block_addr is not None and location.stmt_idx is not None
        key = (
            (location.block_addr, location.stmt_idx)
            if location.block_idx is None
            else (location.block_addr, location.block_idx, location.stmt_idx)
        )
        if overwrite:
            if location.ins_addr is not None:
                self._insn_to_variable[location.ins_addr] = {var_and_offset}
            self._stmt_to_variable[key] = {var_and_offset}
            self._variable_to_stmt[variable].add(key)
            if atom_hash is not None:
                self._atom_to_variable[key][atom_hash] = {var_and_offset}
            if isinstance(atom, ailment.Expr.VirtualVariable):
                self._vvarid_to_variable[atom.varid] = variable
                self._variable_to_vvarids[variable] = {atom.varid}
        else:
            if location.ins_addr is not None:
                self._insn_to_variable[location.ins_addr].add(var_and_offset)
            self._stmt_to_variable[key].add(var_and_offset)
            self._variable_to_stmt[variable].add(key)
            if atom_hash is not None:
                self._atom_to_variable[key][atom_hash].add(var_and_offset)
            if isinstance(atom, ailment.Expr.VirtualVariable):
                self._vvarid_to_variable[atom.varid] = variable
                self._variable_to_vvarids[variable].add(atom.varid)

    def remove_variable_by_atom(self, location: CodeLocation, variable: SimVariable, atom):
        assert location.block_addr is not None and location.stmt_idx is not None
        key = (
            (location.block_addr, location.stmt_idx)
            if location.block_idx is None
            else (location.block_addr, location.block_idx, location.stmt_idx)
        )
        if key in self._stmt_to_variable:
            for var_and_offset in list(self._stmt_to_variable[key]):
                if var_and_offset[0] == variable:
                    self._stmt_to_variable[key].remove(var_and_offset)
            if not self._stmt_to_variable[key]:
                del self._stmt_to_variable[key]

        atom_hash = (hash(atom) & 0xFFFF_FFFF) if atom is not None else None
        if key in self._atom_to_variable and atom_hash is not None and atom_hash in self._atom_to_variable[key]:
            for var_and_offset in list(self._atom_to_variable[key][atom_hash]):
                if var_and_offset[0] == variable:
                    self._atom_to_variable[key][atom_hash].discard(var_and_offset)
            if not self._atom_to_variable[key][atom_hash]:
                del self._atom_to_variable[key][atom_hash]
            if not self._atom_to_variable[key]:
                del self._atom_to_variable[key]

    def make_phi_node(self, block_addr, *variables):
        """
        Create a phi variable for variables at block `block_addr`.

        :param int block_addr:  The address of the current block.
        :param variables:       Variables that the phi variable represents.
        :return:                The created phi variable.
        """

        existing_phis = set()
        non_phis = set()
        for var in variables:
            if self.is_phi_variable(var):
                existing_phis.add(var)
            else:
                non_phis.add(var)
            if var in self._variables_to_phivars:
                for phivar in self._variables_to_phivars[var]:
                    existing_phis.add(phivar)

        if len(existing_phis) >= 1:
            # iterate through existing phi variables to see if any of it is already used as the phi variable for this
            # block. if so, we reuse it to avoid redundant variable allocations
            for phi in existing_phis:
                if block_addr in self._phi_variables_by_block and phi in self._phi_variables_by_block[block_addr]:
                    if not non_phis.issubset(self.get_phi_subvariables(phi)):
                        # Update the variables that this phi variable represents
                        self._phi_variables[phi] |= non_phis
                    return phi

        # allocate a new phi variable
        repre = sorted(variables, key=lambda val: val.key)[0]
        repre_type = type(repre)
        repre_size = max(var.size for var in variables)
        if repre_type is SimRegisterVariable:
            ident_sort = "register"
            a = SimRegisterVariable(repre.reg, repre_size, ident=self.next_variable_ident(ident_sort))
        elif repre_type is SimMemoryVariable:
            ident_sort = "global"
            a = SimMemoryVariable(repre.addr, repre_size, ident=self.next_variable_ident(ident_sort))
        elif repre_type is SimStackVariable:
            ident_sort = "stack"
            a = SimStackVariable(repre.offset, repre_size, ident=self.next_variable_ident(ident_sort))
        else:
            raise TypeError(f'make_phi_node(): Unsupported variable type "{type(repre)}".')

        # Keep a record of all phi variables
        self._phi_variables[a] = set(variables)
        self._phi_variables_by_block[block_addr].add(a)
        for var in variables:
            self._variables_to_phivars[var].add(a)

        return a

    def set_live_variables(self, addr, register_region, stack_region):
        lv = LiveVariables(register_region, stack_region)
        self._live_variables[addr] = lv

    def find_variables_by_insn(self, ins_addr, sort):
        if ins_addr not in self._insn_to_variable:
            return None

        if sort in (VariableType.MEMORY, "memory"):
            vars_and_offset = [
                (var, offset)
                for var, offset in self._insn_to_variable[ins_addr]
                if isinstance(var, (SimStackVariable, SimMemoryVariable))
            ]
        elif sort in (VariableType.REGISTER, "register"):
            vars_and_offset = [
                (var, offset)
                for var, offset in self._insn_to_variable[ins_addr]
                if isinstance(var, SimRegisterVariable)
            ]
        else:
            l.error('find_variable_by_insn(): Unsupported variable sort "%s".', sort)
            return []

        return vars_and_offset

    def is_variable_used_at(self, variable: SimVariable, loc: tuple[int, int]) -> bool:
        return loc in self._variable_to_stmt[variable]

    def find_variable_by_stmt(self, block_addr, stmt_idx, sort, block_idx: int | None = None):
        variables = sorted(
            self.find_variables_by_stmt(block_addr, stmt_idx, sort, block_idx=block_idx),
            key=lambda var: (var[1], var[0].key),
        )
        if variables:
            return variables[0]
        return None

    def find_variables_by_stmt(
        self, block_addr: int, stmt_idx: int, sort: str, block_idx: int | None = None
    ) -> list[tuple[SimVariable, int]]:
        key = (block_addr, stmt_idx) if block_idx is None else (block_addr, block_idx, stmt_idx)

        if key not in self._stmt_to_variable:
            return []

        variables = self._stmt_to_variable[key]
        if not variables:
            return []

        var_and_offsets: list[tuple[SimVariable, int]]
        if sort == "memory":
            var_and_offsets = [
                (var, offset)
                for var, offset in self._stmt_to_variable[key]
                if isinstance(var, (SimStackVariable, SimMemoryVariable))
            ]
        elif sort == "register":
            var_and_offsets = [
                (var, offset) for var, offset in self._stmt_to_variable[key] if isinstance(var, SimRegisterVariable)
            ]
        else:
            l.error('find_variables_by_stmt(): Unsupported variable sort "%s".', sort)
            return []

        return var_and_offsets

    def find_variable_by_atom(self, block_addr, stmt_idx, atom, block_idx: int | None = None):
        variables = sorted(
            self.find_variables_by_atom(block_addr, stmt_idx, atom, block_idx=block_idx),
            key=lambda val: (val[1], val[0].key),
        )
        if variables:
            return variables[0]
        return None

    def find_variables_by_atom(
        self, block_addr, stmt_idx, atom, block_idx: int | None = None
    ) -> set[tuple[SimVariable, int]]:
        key = (block_addr, stmt_idx) if block_idx is None else (block_addr, block_idx, stmt_idx)

        if key not in self._atom_to_variable:
            return set()

        atom_hash = hash(atom) & 0xFFFF_FFFF
        if atom_hash not in self._atom_to_variable[key]:
            return set()

        return self._atom_to_variable[key][atom_hash]

    def find_variables_by_stack_offset(self, offset: int) -> set[SimVariable]:
        return self._stack_region.get_variables_by_offset(offset)

    def find_variables_by_register(self, reg: str | int) -> set[SimVariable]:
        if type(reg) is str:
            reg = self.manager._kb._project.arch.registers.get(reg)[0]
        return self._register_region.get_variables_by_offset(reg)

    def get_variable_accesses(self, variable: SimVariable, same_name: bool = False) -> list[VariableAccess]:
        if not same_name:
            if variable in self._variable_accesses:
                return list(self._variable_accesses[variable])

            return []

        # find all variables with the same variable name

        vars_list = []

        for var in self._variable_accesses:
            if variable.name == var.name:
                vars_list.append(var)

        accesses: list[VariableAccess] = []
        for var in vars_list:
            accesses.extend(self.get_variable_accesses(var))

        return accesses

    @overload
    def get_variables(self, sort: Literal["stack"], collapse_same_ident: bool = False) -> list[SimStackVariable]: ...
    @overload
    def get_variables(self, sort: Literal["reg"], collapse_same_ident: bool = False) -> list[SimRegisterVariable]: ...
    @overload
    def get_variables(
        self, sort: None = None, collapse_same_ident: bool = False
    ) -> list[SimRegisterVariable | SimRegisterVariable]: ...

    def get_variables(self, sort=None, collapse_same_ident=False):
        """
        Get a list of variables.

        :param sort:                Sort of the variable to get.
        :param collapse_same_ident: Whether variables of the same identifier should be collapsed or not.
        :return:                    A list of variables.
        """

        variables = []

        if collapse_same_ident:
            raise NotImplementedError

        for var in self._variables:
            if sort == "stack" and not isinstance(var, SimStackVariable):
                continue
            if sort == "reg" and not isinstance(var, SimRegisterVariable):
                continue
            variables.append(var)

        return variables

    @overload
    def get_unified_variables(self, sort: Literal["stack"]) -> list[SimStackVariable]: ...
    @overload
    def get_unified_variables(self, sort: Literal["reg"]) -> list[SimRegisterVariable]: ...
    @overload
    def get_unified_variables(self, sort: None) -> list[SimRegisterVariable | SimRegisterVariable]: ...

    def get_unified_variables(self, sort=None):
        """
        Get a list of unified variables.

        :param sort:    Sort of the variable to get.
        :return:        A list of variables.
        """

        variables = []

        for var in self._unified_variables:
            if sort == "stack" and not isinstance(var, SimStackVariable):
                continue
            if sort == "reg" and not isinstance(var, SimRegisterVariable):
                continue
            variables.append(var)

        return variables

    def get_global_variables(self, addr):
        """
        Get global variable by the address of the variable.

        :param int addr:    Address of the variable.
        :return:            A set of variables or an empty set if no variable exists.
        """
        return self._global_region.get_variables_by_offset(addr)

    def is_phi_variable(self, var):
        """
        Test if `var` is a phi variable.

        :param SimVariable var: The variable instance.
        :return:                True if `var` is a phi variable, False otherwise.
        :rtype:                 bool
        """

        return var in self._phi_variables

    def get_phi_subvariables(self, var):
        """
        Get sub-variables that phi variable `var` represents.

        :param SimVariable var: The variable instance.
        :return:                A set of sub-variables, or an empty set if `var` is not a phi variable.
        :rtype:                 set
        """

        if not self.is_phi_variable(var):
            return set()
        return self._phi_variables[var]

    def get_phi_variables(self, block_addr):
        """
        Get a dict of phi variables and their corresponding variables.

        :param int block_addr:  Address of the block.
        :return:                A dict of phi variables of an empty dict if there are no phi variables at the block.
        :rtype:                 dict
        """

        if block_addr not in self._phi_variables_by_block:
            return {}
        variables = {}
