# pylint:disable=missing-class-docstring,too-many-boolean-expressions
from __future__ import annotations
import enum
from collections import defaultdict
from contextlib import suppress
import logging

import networkx
from sortedcontainers import SortedDict

from angr.utils.constants import MAX_POINTSTO_BITS
from .typevars import (
    Existence,
    Subtype,
    Equivalence,
    Add,
    Sub,
    TypeVariable,
    DerivedTypeVariable,
    HasField,
    AddN,
    SubN,
    IsArray,
    TypeConstraint,
    Load,
    Store,
    BaseLabel,
    FuncIn,
    FuncOut,
    ConvertTo,
    new_dtv,
)
from .typeconsts import (
    BottomType,
    TopType,
    TypeConstant,
    Int,
    Int8,
    Int16,
    Int32,
    Int64,
    Int128,
    Int256,
    Int512,
    Pointer,
    Pointer32,
    Pointer64,
    Struct,
    Array,
    Function,
    int_type,
    Float,
    Float32,
    Float64,
)
from .variance import Variance
from .dfa import DFAConstraintSolver, EmptyEpsilonNFAError

_l = logging.getLogger(__name__)


Top_ = TopType()
Int_ = Int()
Int512_ = Int512()
Int256_ = Int256()
Int128_ = Int128()
Int64_ = Int64()
Int32_ = Int32()
Int16_ = Int16()
Int8_ = Int8()
Bottom_ = BottomType()
Pointer64_ = Pointer64()
Pointer32_ = Pointer32()
Struct_ = Struct()
Array_ = Array()
Float_ = Float()
Float32_ = Float32()
Float64_ = Float64()


PRIMITIVE_TYPES = {
    Top_,
    Int_,
    Int8_,
    Int16_,
    Int32_,
    Int64_,
    Int128_,
    Int256_,
    Int512_,
    Pointer32_,
    Pointer64_,
    Bottom_,
    Struct_,
    Array_,
    Float_,
    Float32_,
    Float64_,
}


# lattice for 64-bit binaries
BASE_LATTICE_64 = networkx.DiGraph()
BASE_LATTICE_64.add_edge(Top_, Int_)
BASE_LATTICE_64.add_edge(Int_, Int512_)
BASE_LATTICE_64.add_edge(Int_, Int256_)
BASE_LATTICE_64.add_edge(Int_, Int128_)
BASE_LATTICE_64.add_edge(Int_, Int64_)
BASE_LATTICE_64.add_edge(Int_, Int32_)
BASE_LATTICE_64.add_edge(Int_, Int16_)
BASE_LATTICE_64.add_edge(Int_, Int8_)
BASE_LATTICE_64.add_edge(Int32_, Bottom_)
BASE_LATTICE_64.add_edge(Int16_, Bottom_)
BASE_LATTICE_64.add_edge(Int8_, Bottom_)
BASE_LATTICE_64.add_edge(Int64_, Pointer64_)
BASE_LATTICE_64.add_edge(Pointer64_, Bottom_)

# lattice for 32-bit binaries
BASE_LATTICE_32 = networkx.DiGraph()
BASE_LATTICE_32.add_edge(Top_, Int_)
BASE_LATTICE_32.add_edge(Int_, Int512_)
BASE_LATTICE_32.add_edge(Int_, Int256_)
BASE_LATTICE_32.add_edge(Int_, Int128_)
BASE_LATTICE_32.add_edge(Int_, Int64_)
BASE_LATTICE_32.add_edge(Int_, Int32_)
BASE_LATTICE_32.add_edge(Int_, Int16_)
BASE_LATTICE_32.add_edge(Int_, Int8_)
BASE_LATTICE_32.add_edge(Int32_, Pointer32_)
BASE_LATTICE_32.add_edge(Int64_, Bottom_)
BASE_LATTICE_32.add_edge(Pointer32_, Bottom_)
BASE_LATTICE_32.add_edge(Int16_, Bottom_)
BASE_LATTICE_32.add_edge(Int8_, Bottom_)

BASE_LATTICES = {
    32: BASE_LATTICE_32,
    64: BASE_LATTICE_64,
}


#
# Sketch
#


class SketchNodeBase:
    """
    The base class for nodes in a sketch.
    """

    __slots__ = ()


class SketchNode(SketchNodeBase):
    """
    Represents a node in a sketch graph.
    """

    __slots__ = ("lower_bound", "typevar", "upper_bound")

    def __init__(self, typevar: TypeVariable | DerivedTypeVariable):
        self.typevar: TypeVariable | DerivedTypeVariable = typevar
        self.upper_bound: TypeConstant = TopType()
        self.lower_bound: TypeConstant = BottomType()

    def __repr__(self):
        return f"{self.lower_bound} <: {self.typevar} <: {self.upper_bound}"

    def __eq__(self, other):
        return isinstance(other, SketchNode) and self.typevar == other.typevar

    def __hash__(self):
        return hash((SketchNode, self.typevar))

    @property
    def size(self) -> int | None:
        """
        Best-effort estimation of the size of the typevar (in bits). Returns None if we cannot determine.
        """

        if isinstance(self.typevar, DerivedTypeVariable):
            last_label = self.typevar.labels[-1]
            if isinstance(last_label, HasField) and last_label.bits != MAX_POINTSTO_BITS:
                return last_label.bits
        if isinstance(self.lower_bound, TypeConstant) and not isinstance(self.lower_bound, (TopType, BottomType)):
            with suppress(NotImplementedError):
                return self.lower_bound.size * 8
        if isinstance(self.upper_bound, TypeConstant) and not isinstance(self.upper_bound, (TopType, BottomType)):
            with suppress(NotImplementedError):
                return self.upper_bound.size * 8
        return None


class RecursiveRefNode(SketchNodeBase):
    """
    Represents a cycle in a sketch graph.

    This is equivalent to sketches.LabelNode in the reference implementation of retypd.
    """

    def __init__(self, target: TypeVariable | DerivedTypeVariable):
        self.target: TypeVariable | DerivedTypeVariable = target

    def __repr__(self):
        return f"Ref({self.target})"

    def __hash__(self):
        return hash((RecursiveRefNode, self.target))

    def __eq__(self, other):
        return type(other) is RecursiveRefNode and other.target == self.target


class Sketch:
    """
    Describes the sketch of a type variable.
    """

    __slots__ = (
        "graph",
        "node_mapping",
        "root",
        "solver",
    )

    def __init__(self, solver: SimpleSolver, root: TypeVariable):
        self.root: SketchNode = SketchNode(root)
        self.graph = networkx.MultiDiGraph()
        self.node_mapping: dict[TypeVariable | DerivedTypeVariable, SketchNodeBase] = {}
        self.solver = solver

        # add the root node
        self.graph.add_node(self.root)
        self.node_mapping[root] = self.root

    def lookup(self, typevar: TypeVariable | DerivedTypeVariable) -> SketchNodeBase | None:
        if typevar in self.node_mapping:
            return self.node_mapping[typevar]
        node: SketchNodeBase | None = None
        if isinstance(typevar, DerivedTypeVariable):
            t = SimpleSolver._to_typevar_or_typeconst(typevar.type_var)
            assert isinstance(t, TypeVariable)
            node = self.node_mapping[t]
            for label in typevar.labels:
                succs = []
                for _, dst, data in self.graph.out_edges(node, data=True):
                    if "label" in data and data["label"] == label and dst not in succs:
                        succs.append(dst)
                if len(succs) > 1:
                    _l.warning(
                        "Multiple successors found for node %s with label %s. Picking the first one.", node, label
                    )
                    succs = succs[:1]
                if not succs:
                    return None
                node = succs[0]
                if isinstance(node, RecursiveRefNode):
                    node = self.lookup(node.target)
        return node

    def add_edge(self, src: SketchNodeBase, dst: SketchNodeBase, label) -> None:
        # ensure the label does not already exist in existing edges
        if self.graph.has_edge(src, dst):
            for data in self.graph.get_edge_data(src, dst).values():
                if "label" in data and data["label"] == label:
                    return
        if isinstance(src, SketchNode) and src.typevar not in self.node_mapping:
            self.node_mapping[src.typevar] = src
        if isinstance(dst, SketchNode) and dst.typevar not in self.node_mapping:
            self.node_mapping[dst.typevar] = dst
        self.graph.add_edge(src, dst, label=label)

    def add_constraint(self, constraint: TypeConstraint) -> None:
        # sub <: super
        if not isinstance(constraint, Subtype):
            return
        subtype, _ = self.flatten_typevar(constraint.sub_type)
        supertype, try_maxsize = self.flatten_typevar(constraint.super_type)

        if (
            try_maxsize
            and isinstance(subtype, TypeVariable)
            and subtype in self.solver.stackvar_max_sizes
            and isinstance(supertype, TypeConstant)
            and not isinstance(supertype, BottomType)
        ):
            basetype = supertype
            if not isinstance(basetype, (TopType, BottomType)):
                assert basetype.size is not None
                max_size = self.solver.stackvar_max_sizes.get(subtype, None)
                if max_size not in {0, None} and basetype.size > 0 and max_size // basetype.size > 0:  # type: ignore
                    supertype = Array(element=basetype, count=max_size // basetype.size)  # type: ignore

        if SimpleSolver._typevar_inside_set(subtype, PRIMITIVE_TYPES) and not SimpleSolver._typevar_inside_set(
            supertype, PRIMITIVE_TYPES
        ):
            assert isinstance(supertype, (TypeVariable, DerivedTypeVariable))
            super_node = self.lookup(supertype)
            assert super_node is None or isinstance(super_node, SketchNode)
            if super_node is not None:
                super_node.lower_bound = self.solver.join(super_node.lower_bound, subtype)
        elif SimpleSolver._typevar_inside_set(supertype, PRIMITIVE_TYPES) and not SimpleSolver._typevar_inside_set(
            subtype, PRIMITIVE_TYPES
        ):
            assert isinstance(subtype, (TypeVariable, DerivedTypeVariable))
            sub_node = self.lookup(subtype)
            assert sub_node is None or isinstance(sub_node, SketchNode)
            # assert sub_node is not None
            if sub_node is not None:
                sub_node.upper_bound = self.solver.meet(sub_node.upper_bound, supertype)

    @staticmethod
    def flatten_typevar(
        derived_typevar: TypeVariable | TypeConstant | DerivedTypeVariable,
    ) -> tuple[DerivedTypeVariable | TypeVariable | TypeConstant, bool]:
        # pylint:disable=too-many-boolean-expressions
        if (
            isinstance(derived_typevar, DerivedTypeVariable)
            and isinstance(derived_typevar.type_var, Pointer)
            and SimpleSolver._typevar_inside_set(derived_typevar.type_var.basetype, PRIMITIVE_TYPES)
            and len(derived_typevar.labels) == 2
            and isinstance(derived_typevar.labels[0], Load)
            and isinstance(derived_typevar.labels[1], HasField)
            and derived_typevar.labels[1].offset == 0
            and derived_typevar.labels[1].bits == MAX_POINTSTO_BITS
        ):
            bt = derived_typevar.type_var.basetype
            assert bt is not None
            return bt, True
        return derived_typevar, False


#
# Constraint graph
#


class ConstraintGraphTag(enum.Enum):
    LEFT = 0
    RIGHT = 1
    UNKNOWN = 2


class FORGOTTEN(enum.Enum):
    PRE_FORGOTTEN = 0
    POST_FORGOTTEN = 1


class ConstraintGraphNode:
    __slots__ = ("forgotten", "tag", "typevar", "variance")

    def __init__(
        self,
        typevar: TypeVariable | DerivedTypeVariable | TypeConstant,
        variance: Variance,
        tag: ConstraintGraphTag,
        forgotten: FORGOTTEN,
    ):
        self.typevar = typevar
        self.variance = variance
        self.tag = tag
        self.forgotten = forgotten

    def __repr__(self):
        variance_str = "CO" if self.variance == Variance.COVARIANT else "CONTRA"
        if self.tag == ConstraintGraphTag.LEFT:
            tag_str = "L"
        elif self.tag == ConstraintGraphTag.RIGHT:
            tag_str = "R"
        else:
            tag_str = "U"
        forgotten_str = "PRE" if self.forgotten == FORGOTTEN.PRE_FORGOTTEN else "POST"
        s = f"{self.typevar}#{variance_str}.{tag_str}.{forgotten_str}"
        if ":" in s:
            return '"' + s + '"'
        return s

    def __eq__(self, other):
        if not isinstance(other, ConstraintGraphNode):
            return False
        return (
            self.typevar == other.typevar
            and self.variance == other.variance
            and self.tag == other.tag
            and self.forgotten == other.forgotten
        )

    def __hash__(self):
        return hash((ConstraintGraphNode, self.typevar, self.variance, self.tag, self.forgotten))

    def forget_last_label(self) -> tuple[ConstraintGraphNode, BaseLabel] | None:
        if isinstance(self.typevar, DerivedTypeVariable) and self.typevar.labels:
            last_label = self.typevar.labels[-1]
            if len(self.typevar.labels) == 1:
                prefix = self.typevar.type_var
            else:
                prefix = DerivedTypeVariable(self.typevar.type_var, None, labels=self.typevar.labels[:-1])
            variance = Variance.COVARIANT if self.variance == last_label.variance else Variance.CONTRAVARIANT
            if not isinstance(prefix, (TypeVariable, DerivedTypeVariable)):
                # we may see incorrectly generated type constraints that attempt to load from an int:
                #   int64.load
                # we don't want to entertain such constraints
                return None
            return (
                ConstraintGraphNode(prefix, variance, self.tag, FORGOTTEN.PRE_FORGOTTEN),
                self.typevar.labels[-1],
            )
        return None

    def recall(self, label: BaseLabel) -> ConstraintGraphNode:
        if isinstance(self.typevar, DerivedTypeVariable):
            labels = (*self.typevar.labels, label)
            typevar = self.typevar.type_var
        elif isinstance(self.typevar, (TypeVariable, TypeConstant)):
            labels = (label,)
            typevar = self.typevar
        else:
            raise TypeError(f"Unsupported type {type(self.typevar)}")
        variance = Variance.COVARIANT if self.variance == label.variance else Variance.CONTRAVARIANT
        var = typevar if not labels else DerivedTypeVariable(typevar, None, labels=labels)
        assert isinstance(var, (TypeVariable, DerivedTypeVariable))
        return ConstraintGraphNode(var, variance, self.tag, FORGOTTEN.PRE_FORGOTTEN)

    def inverse(self) -> ConstraintGraphNode:
        if self.tag == ConstraintGraphTag.LEFT:
            tag = ConstraintGraphTag.RIGHT
        elif self.tag == ConstraintGraphTag.RIGHT:
            tag = ConstraintGraphTag.LEFT
        else:
            tag = ConstraintGraphTag.UNKNOWN

        variance = Variance.CONTRAVARIANT if self.variance == Variance.COVARIANT else Variance.COVARIANT

        return ConstraintGraphNode(self.typevar, variance, tag, self.forgotten)

    def inverse_wo_tag(self) -> ConstraintGraphNode:
        """
        Invert the variance only.
        """
        variance = Variance.CONTRAVARIANT if self.variance == Variance.COVARIANT else Variance.COVARIANT

        return ConstraintGraphNode(self.typevar, variance, self.tag, self.forgotten)


#
# The solver
#


class SimpleSolver:
    """
    SimpleSolver is, by its name, a simple solver. Most of this solver is based on the (complex) simplification logic
    that the retypd paper describes and the retypd re-implementation (https://github.com/GrammaTech/retypd) implements.
    Additionally, we add some improvements to allow type propagation of known struct names, among a few other
    improvements.
    """

    def __init__(
        self,
        bits: int,
        constraints,
        typevars: dict[TypeVariable, set[TypeVariable]],
        constraint_set_degradation_threshold: int = 150,
        stackvar_max_sizes: dict[TypeVariable, int] | None = None,
    ):
        if bits not in (32, 64):
            raise ValueError(f"Pointer size {bits} is not supported. Expect 32 or 64.")

        self.bits = bits
        self._constraints: dict[TypeVariable, set[TypeConstraint]] = constraints
        self._typevars: dict[TypeVariable, set[TypeVariable]] = typevars
        self.stackvar_max_sizes = stackvar_max_sizes if stackvar_max_sizes is not None else {}
        self._constraint_set_degradation_threshold = constraint_set_degradation_threshold
        self._base_lattice = BASE_LATTICES[bits]
        self._base_lattice_inverted = networkx.DiGraph()
        for src, dst in self._base_lattice.edges:
            self._base_lattice_inverted.add_edge(dst, src)

        # statistics
        self.processed_constraints_count: int = 0
        self.simplified_constraints_count: int = 0
        self.eqclass_constraints_count: list[int] = []

        #
        # Solving state
        #
        self._equivalence = defaultdict(dict)
        for func_tv in list(self._constraints):
            if self._constraints[func_tv]:
                self.processed_constraints_count += len(self._constraints[func_tv])
                self.preprocess(func_tv)
                self.simplified_constraints_count += len(self._constraints[func_tv])

        self.solution = {}
        for tv, sol in self._equivalence.items():
            if isinstance(tv, TypeVariable) and isinstance(sol, TypeConstant):
                self.solution[tv] = sol

        self._solution_cache = {}
        self.solve()
        for func_tv in list(self._constraints):
            self._convert_arrays(self._constraints[func_tv])

    def preprocess(self, func_tv: TypeVariable):
        self._constraints[func_tv] |= self._eq_constraints_from_tvs(self._constraints[func_tv])
        ptr_tvs = self._ptr_tvs_from_constraints(self._constraints[func_tv])
        self._constraints[func_tv] = self._remove_alignment_int_ptr_subtyping_constraints(
            self._constraints[func_tv], ptr_tvs
        )
        self._constraints[func_tv] |= self._eq_constraints_from_add(func_tv, ptr_tvs)
        self._constraints[func_tv] |= self._eq_constraints_from_sub(func_tv, ptr_tvs)
        self._constraints[func_tv] |= self._discover_equivalence(self._constraints[func_tv])
        new_constraints, replacements = self._handle_equivalence(self._constraints[func_tv])
        self._equivalence |= replacements
        self._constraints[func_tv] = new_constraints
        self._constraints[func_tv] = self._filter_constraints(self._constraints[func_tv])

    def solve(self):
        """
        Steps:

        For each type variable,
        - Infer the shape in its sketch
        - Build the constraint graph
        - Collect all constraints
        - Apply constraints to derive the lower and upper bounds
        - Determine a solution for type variables with constraints
        - Rewrite the constraint graph by replacing determined type variables with their solutions
        - Solve repeatedly until all interesting type variables have solutions

        By repeatedly solving until exhausting interesting type variables, we ensure the S-Trans rule is applied.
        """

        prem_typevars = set(self._constraints) | next(iter(self._typevars.values()))
        typevars = set()
        for tv in prem_typevars:
            if tv not in self._equivalence:
                typevars.add(tv)
            else:
                repl = self._equivalence[tv]
                if isinstance(repl, TypeVariable):
                    typevars.add(repl)

        constraints = set()

        for tv in typevars:
            if tv in self._constraints:
                constraints |= self._constraints[tv]

        equiv_classes, prelim_sketches = self.infer_shapes(typevars, constraints)

        # only create sketches for the type variables representing their equivalence classes
        tv_to_reptvs = {}
        for tv_or_dtv, reptv in equiv_classes.items():
            if not isinstance(tv_or_dtv, DerivedTypeVariable) and tv_or_dtv is not reptv:
                tv_to_reptvs[tv_or_dtv] = reptv
        # rewrite constraints to only use representative type variables
        constraints = self._rewrite_constraints_with_replacements(constraints, tv_to_reptvs)
        # update sketches
        sketches = {tv_to_reptvs.get(tv, tv): sketch for tv, sketch in prelim_sketches.items()}
        # rewrite typevars as well...
        typevars = {tv_to_reptvs.get(tv, tv) for tv in typevars}

        # collect typevars used in the constraint set
        constrained_typevars = set()
        for constraint in constraints:
            if isinstance(constraint, Subtype):
                for t in (constraint.sub_type, constraint.super_type):
                    if isinstance(t, DerivedTypeVariable):
                        if t.type_var in typevars:
                            constrained_typevars.add(t.type_var)
                    elif isinstance(t, TypeVariable) and t in typevars:
                        constrained_typevars.add(t)

        constraintset2tvs = defaultdict(set)
        tvs_seen = set()
        for idx, tv in enumerate(sorted(constrained_typevars, key=lambda x: x.idx)):
            _l.debug("Collecting constraints for type variable %r (%d/%d)", tv, idx + 1, len(constrained_typevars))
            if tv in tvs_seen:
                continue
            # build a sub constraint set for the type variable
            constraint_subset, related_tvs = self._generate_constraint_subset(constraints, {tv})
            # drop all type vars outside constrained_typevars
            related_tvs = related_tvs.intersection(constrained_typevars)
            tvs_seen |= related_tvs
            frozen_constraint_subset = frozenset(constraint_subset)
            constraintset2tvs[frozen_constraint_subset] = related_tvs

        for idx, (constraint_subset, tvs) in enumerate(constraintset2tvs.items()):
            _l.debug(
                "Solving %d constraints for %d type variables %r (%d/%d)",
                len(constraint_subset),
                len(tvs),
                tvs,
                idx + 1,
                len(constraintset2tvs),
            )
            self.eqclass_constraints_count.append(len(constraint_subset))

            constraint_subset = self._drop_missized_constraints(tvs, constraint_subset)

            if len(constraint_subset) > self._constraint_set_degradation_threshold:
                _l.debug(
                    "Constraint subset contains %d constraints, which is over the limit of %d. Enter degradation.",
                    len(constraint_subset),
                    self._constraint_set_degradation_threshold,
                )
                constraint_subset = self._degrade_constraint_set(constraint_subset)
                _l.debug("Degraded constraint subset to %d constraints.", len(constraint_subset))

            while constraint_subset:

                _l.debug("Working with %d constraints.", len(constraint_subset))

                # remove constraints that are a <: b where a only appears once; in this case, the solution fo a is
                # entirely determined by the solution of b (which is the upper bound of a)
                filtered_constraint_subset, ub_subtypes = self._filter_leaf_typevars(constraint_subset, tvs)
                _l.debug(
                    "Filtered %d leaf typevars; %d constraints remain.",
                    len(ub_subtypes),
                    len(filtered_constraint_subset),
                )

                base_constraint_graph = self._generate_constraint_graph(
                    filtered_constraint_subset, tvs | PRIMITIVE_TYPES
                )
                primitive_constraints = self._generate_primitive_constraints(tvs, base_constraint_graph)
                if len(tvs) > 1:
                    primitive_constraints |= self._generate_transitive_subtype_constraints(
                        tvs, filtered_constraint_subset, primitive_constraints
                    )
                tvs_with_primitive_constraints = set()
                for primitive_constraint in primitive_constraints:
                    tv = self._typevar_from_primitive_constraint(primitive_constraint)
                    tvs_with_primitive_constraints.add(tv)
                    assert tv is not None, f"Cannot find type variable in primitive constraint {primitive_constraint}"
                    sketches[tv].add_constraint(primitive_constraint)
                solutions = {}
                self.determine(
                    sketches, sorted(tvs_with_primitive_constraints, key=lambda x: x.idx), equiv_classes, solutions
                )
                _l.debug("Determined solutions for %d type variable(s).", len(tvs_with_primitive_constraints))

                leaf_solutions = 0
                for tv_, ub_tv in ub_subtypes.items():
                    if ub_tv in solutions:
                        solutions[tv_] = solutions[ub_tv]
                        leaf_solutions += 1
                    elif isinstance(ub_tv, TypeConstant):
                        solutions[tv_] = ub_tv
                        leaf_solutions += 1
                _l.debug("Determined solutions for %d leaf type variable(s).", leaf_solutions)

                if not solutions:
                    break
                self.solution |= solutions

                tvs = {tv for tv in tvs if tv not in solutions}
                if not tvs:
                    break
                # rewrite existing constraints
                new_constraint_subset = set()
                for constraint in constraint_subset:
                    rewritten = self._rewrite_constraint(constraint, solutions)
                    new_constraint_subset.add(rewritten)
                constraint_subset = self._filter_constraints(new_constraint_subset)

        # set the solution for missing type vars to TOP
        self.determine(sketches, set(sketches).difference(set(self.solution)), equiv_classes, self.solution)

        # set solutions for non-representative type variables
        for tv, reptv in equiv_classes.items():
            if reptv in self.solution and tv not in self.solution:
                self.solution[tv] = self.solution[reptv]

    def infer_shapes(
        self, typevars: set[TypeVariable], constraints: set[TypeConstraint]
    ) -> tuple[dict, dict[TypeVariable, Sketch]]:
        """
        Computing sketches from constraint sets. Implements Algorithm E.1 in the retypd paper.
        """

        equivalence_classes, quotient_graph = self.compute_quotient_graph(constraints)

        sketches: dict[TypeVariable, Sketch] = {}
        for tv in typevars:
            rep_tv = equivalence_classes.get(tv, tv)
            if rep_tv not in sketches:
                sketches[rep_tv] = Sketch(self, rep_tv)

        for tv, sketch in sketches.items():
            sketch_node = sketch.lookup(tv)
            graph_node = equivalence_classes.get(tv, None)
            # assert graph_node is not None
            if graph_node is None or sketch_node is None:
                continue
            assert isinstance(graph_node, TypeVariable)
            assert isinstance(sketch_node, SketchNode)
            visited = {graph_node: sketch_node}
            self._get_all_paths(quotient_graph, sketch, graph_node, visited)
        return equivalence_classes, sketches

    def compute_quotient_graph(self, constraints: set[TypeConstraint]):
        """
        Compute the quotient graph (the constraint graph modulo ~ in Algorithm E.1 in the retypd paper) with respect to
        a given set of type constraints.
        """

        g = networkx.DiGraph()
        # collect all derived type variables
        typevars = self._typevars_from_constraints(constraints)
        g.add_nodes_from(typevars)
        # add paths for each derived type variable into the graph
        for tv in typevars:
            last_node = tv
            prefix = tv
            while isinstance(prefix, DerivedTypeVariable) and prefix.labels:
                assert isinstance(last_node, DerivedTypeVariable)
                prefix = prefix.longest_prefix()
                if prefix is None:
                    continue
                g.add_edge(prefix, last_node, label=last_node.labels[-1])
                last_node = prefix

        # compute the constraint graph modulo ~
        equivalence_classes = {node: node for node in g}

        load = Load()
        store = Store()
        for node in g.nodes:
            lbl_to_node = {}
            for succ in g.successors(node):
                lbl_to_node[succ.labels[-1]] = succ
                if load in lbl_to_node and store in lbl_to_node:
                    self._unify(equivalence_classes, lbl_to_node[load], lbl_to_node[store], g)

        for constraint in constraints:
            if isinstance(constraint, Subtype):
                if self._typevar_inside_set(constraint.super_type, PRIMITIVE_TYPES) or self._typevar_inside_set(
                    constraint.sub_type, PRIMITIVE_TYPES
                ):
                    continue
                self._unify(equivalence_classes, constraint.super_type, constraint.sub_type, g)

        out_graph = networkx.MultiDiGraph()  # there can be multiple edges between two nodes, each edge is associated
        # with a different label
        for src, dst, data in g.edges(data=True):
            src_cls = equivalence_classes.get(src, src)
            dst_cls = equivalence_classes.get(dst, dst)
            label = None if not data else data["label"]
            if label is not None and out_graph.has_edge(src_cls, dst_cls):
                # do not add the same edge twice
                existing_labels = {
                    data_["label"]
                    for _, dst_cls_, data_ in out_graph.out_edges(src_cls, data=True)
                    if dst_cls_ == dst_cls and data
                }
                if label in existing_labels:
                    continue
            out_graph.add_edge(src_cls, dst_cls, label=label)

        return equivalence_classes, out_graph

    def _generate_primitive_constraints(
        self,
        non_primitive_endpoints: set[TypeVariable | DerivedTypeVariable],
        constraint_graph,
    ) -> set[TypeConstraint]:
        constraints_0 = self._solve_constraints_between(constraint_graph, non_primitive_endpoints, PRIMITIVE_TYPES)
        constraints_1 = self._solve_constraints_between(constraint_graph, PRIMITIVE_TYPES, non_primitive_endpoints)
        return constraints_0 | constraints_1

    @staticmethod
    def _generate_transitive_subtype_constraints(
        typevars: set[TypeVariable | DerivedTypeVariable],
        constraints: set[TypeConstraint],
        primitive_constraints: set[TypeConstraint],
    ) -> set[TypeConstraint]:
        """
        Handling multiple type variables at once means we may miss some subtyping relationships between a type variable
        and a primitive type (e.g., tv_1 <: tv_2 and tv_2 <: int, we may not see tv_1 <: int). This method attempts to
        recover some of these missing constraints.
        """

        tv_supertype = defaultdict(set)
        for constraint in primitive_constraints:
            if isinstance(constraint, Subtype) and SimpleSolver._typevar_inside_set(
                constraint.super_type, PRIMITIVE_TYPES
            ):
                tv_supertype[constraint.sub_type].add(constraint.super_type)

        additional_constraints = set()
        for constraint in constraints:
            if (
                isinstance(constraint, Subtype)
                and isinstance(constraint.sub_type, TypeVariable)
                and (
                    (isinstance(constraint.sub_type, DerivedTypeVariable) and constraint.sub_type.type_var in typevars)
                    or (not isinstance(constraint.sub_type, DerivedTypeVariable) and constraint.sub_type in typevars)
                )
                and constraint.super_type in tv_supertype
            ):
                for supertype in tv_supertype[constraint.super_type]:
                    additional_constraints.add(Subtype(constraint.sub_type, supertype))

        return additional_constraints

    @staticmethod
    def _typevars_from_constraints(constraints: set[TypeConstraint]) -> set[TypeVariable | DerivedTypeVariable]:
        """
        Collect derived type variables from a set of constraints.
        """

        typevars: set[TypeVariable | DerivedTypeVariable] = set()
        for constraint in constraints:
            if isinstance(constraint, Subtype):
                if isinstance(constraint.sub_type, TypeVariable) and not (
                    isinstance(constraint.sub_type, DerivedTypeVariable)
                    and isinstance(constraint.sub_type.type_var, TypeConstant)
                ):
                    typevars.add(constraint.sub_type)
                if isinstance(constraint.super_type, TypeVariable) and not (
                    isinstance(constraint.super_type, DerivedTypeVariable)
                    and isinstance(constraint.super_type.type_var, TypeConstant)
                ):
                    typevars.add(constraint.super_type)
            # TODO: Other types of constraints?
        return typevars

    @staticmethod
    def _typevar_from_primitive_constraint(constraint: TypeConstraint) -> TypeVariable | None:
        if isinstance(constraint, Subtype):
            if (
