from __future__ import annotations

import logging
import operator

import claripy
import claripy.backends.backend_vsa as vsa
from claripy.ast import BV, Base, Bool
from claripy.errors import BackendError, ClaripyBalancerError, ClaripyBalancerUnsatError, ClaripyOperationError
from claripy.operations import commutative_operations, opposites

log = logging.getLogger(__name__)


class Balancer:
    """
    The Balancer is an equation redistributor. The idea is to take an AST and rebalance it to, for example, isolate
    unknown terms on one side of an inequality.
    """

    def __init__(self, c):
        self._truisms = []
        self._ast_hash_map = {}
        self._lower_bounds = {}
        self._upper_bounds = {}

        self.sat = True
        try:
            self._doit(c)
        except ClaripyBalancerUnsatError:
            self.bounds = {}
            self.sat = False
        except BackendError:
            log.debug("Backend error in balancer.", exc_info=True)

    @property
    def compat_ret(self):
        return (self.sat, self.replacements)

    def _replacements_iter(self):
        all_keys = set(self._lower_bounds.keys()) | set(self._upper_bounds.keys())
        for k in all_keys:
            ast = self._ast_hash_map[k]
            max_int = (1 << len(ast)) - 1
            min_int = 0
            mn = self._lower_bounds.get(k, min_int)
            mx = self._upper_bounds.get(k, max_int)
            bound_si = claripy.BVS("bound", len(ast)).annotate(claripy.annotation.StridedIntervalAnnotation(1, mn, mx))
            log.debug("Yielding bound %s for %s.", bound_si, ast)
            if ast.op == "Reverse":
                yield (ast.args[0], ast.intersection(bound_si).reversed)
            else:
                yield (ast, ast.intersection(bound_si))

    def _add_lower_bound(self, o, b):
        if o.hash() in self._lower_bounds:
            old_b = self._lower_bounds[o.hash()]
            b = max(b, old_b)

        self._lower_bounds[o.hash()] = b
        self._ast_hash_map[o.hash()] = o

    def _add_upper_bound(self, o, b):
        if o.hash() in self._upper_bounds:
            old_b = self._upper_bounds[o.hash()]
            b = min(b, old_b)

        self._upper_bounds[o.hash()] = b
        self._ast_hash_map[o.hash()] = o

    @property
    def replacements(self):
        return list(self._replacements_iter())

    #
    # AST helper functions
    #

    @staticmethod
    def _same_bound_bv(a):
        mx = Balancer._max(a)
        mn = Balancer._min(a)
        stride = Balancer._stride(a)
        return claripy.BVS("bounds", len(a)).annotate(claripy.annotation.StridedIntervalAnnotation(stride, mn, mx))

    @staticmethod
    def _cardinality(a):
        return a.cardinality if isinstance(a, Base) else 0

    @staticmethod
    def _min(a, signed=False):
        a = claripy.backends.vsa.simplify(a)
        if a.has_annotation_type(claripy.annotation.RegionAnnotation):
            region_annos = a.get_annotations_by_type(claripy.annotation.RegionAnnotation)
            if len(region_annos) == 1:
                a = next(iter(region_annos)).region_base_addr
            else:
                # unfortunately, this is a real abstract pointer
                # the minimum value will be 0 or MIN_INT
                if signed:
                    return -(1 << (len(a) - 1))
                return 0
        return claripy.backends.vsa.min(a, signed=signed)

    @staticmethod
    def _max(a, signed=False):
        a = claripy.backends.vsa.simplify(a)
        if a.has_annotation_type(claripy.annotation.RegionAnnotation):
            region_annos = a.get_annotations_by_type(claripy.annotation.RegionAnnotation)
            if len(region_annos) == 1:
                a = next(iter(region_annos)).region_base_addr
            else:
                # unfortunately, this is a real abstract pointer
                # the minimum value will be 0 or MIN_INT
                if signed:
                    return (1 << (len(a) - 1)) - 1
                return (1 << len(a)) - 1
        return claripy.backends.vsa.max(a, signed=signed)

    @staticmethod
    def _stride(a):
        a = claripy.backends.vsa.simplify(a)
        si_anno = a.get_annotation(claripy.annotation.StridedIntervalAnnotation)
        if si_anno is not None:
            return si_anno.stride
        return 1

    @staticmethod
    def _range(a, signed=False):
        return (Balancer._min(a, signed=signed), Balancer._max(a, signed=signed))

    #
    # Truism alignment
    #

    @staticmethod
    def _align_truism(truism):
        outer_aligned = Balancer._align_ast(truism)
        inner_aligned = Bool(outer_aligned.op, (Balancer._align_ast(outer_aligned.args[0]), *outer_aligned.args[1:]))

        _, _, inner_aligned_canon = claripy.backends.vsa.simplify(inner_aligned).canonicalize()
        _, _, truism_canon = claripy.backends.vsa.simplify(truism).canonicalize()

        if not inner_aligned_canon.identical(truism_canon):
            log.critical(
                "ERROR: the balancer is messing up an AST. This must be looked into. "
                "Please submit the binary and script to the angr project, if possible. "
                "Outer op is %s and inner op is %s.",
                truism.op,
                truism.args[0].op,
            )
            return truism

        return inner_aligned

    @staticmethod
    def _align_ast(a):
        """
        Aligns the AST so that the argument with the highest cardinality is on the left.

        :return: a new AST.
        """

        try:
            if isinstance(a, BV):
                return Balancer._align_bv(a)
            if isinstance(a, Bool) and len(a.args) == 2 and a.args[1].cardinality > a.args[0].cardinality:
                return Balancer._reverse_comparison(a)
            return a
        except ClaripyBalancerError:
            return a

    @staticmethod
    def _reverse_comparison(a):
        new_op = opposites.get(a.op, None)
        if new_op is None:
            raise ClaripyBalancerError(f"unable to reverse comparison {a.op} (missing from 'opposites')")

        try:
            op = getattr(BV, new_op)
        except AttributeError as err:
            raise ClaripyBalancerError(f"unable to reverse comparison {a.op} (AttributeError)") from err

        try:
            return op(*a.args[::-1])
        except ClaripyOperationError as err:
            raise ClaripyBalancerError(f"unable to reverse comparison {a.op} (ClaripyOperationError)") from err

    @staticmethod
    def _align_bv(a):
        if a.op in commutative_operations:
            return BV(a.op, tuple(sorted(a.args, key=lambda v: -Balancer._cardinality(v))), length=len(a))

        match a.op:
            case "__sub__":
                return Balancer._align_sub(a)
            case _:
                return a

    @staticmethod
    def _align_sub(a):
        cardinalities = [Balancer._cardinality(v) for v in a.args]
        if max(cardinalities) == cardinalities[0]:
            return a

        adjusted = tuple(operator.__neg__(v) for v in a.args[1:]) + a.args[:1]
        return BV("__add__", tuple(sorted(adjusted, key=lambda v: -Balancer._cardinality(v))), length=len(a))

    #
    # Find bounds
    #

    def _doit(self, c):
        """
        This function processes the list of truisms and finds bounds for ASTs.
        """
        self._truisms.append(claripy.excavate_ite(c))

        processed_truisms = set()
        identified_assumptions = set()

        while len(self._truisms):
            truism = self._truisms.pop()

            if truism in processed_truisms:
                continue

            unpacked_truisms = Balancer._unpack_truisms(truism)
            if claripy.backends.vsa.is_false(truism):
                raise ClaripyBalancerUnsatError

            processed_truisms.add(truism)
            if len(unpacked_truisms):
                self._truisms.extend(t for t in unpacked_truisms if not claripy.backends.vsa.is_true(t))
                continue

            if not Balancer._handleable_truism(truism):
                continue

            truism = Balancer._adjust_truism(truism)

            assumptions = Balancer._get_assumptions(truism)
            if truism not in identified_assumptions and len(assumptions):
                log.debug("Queued assumptions %s for truism %s.", assumptions, truism)
                self._truisms.extend(assumptions)
                identified_assumptions.update(assumptions)

            log.debug("Processing truism %s", truism)
            balanced_truism = self._balance(truism)
            log.debug("... handling")
            self._handle(balanced_truism)

    @staticmethod
    def _handleable_truism(t):
        """
        Checks whether we can handle this truism. The truism should already be aligned.
        """
        if len(t.args) < 2:
            log.debug("can't do anything with an unop bool")
            return None
        if t.args[0].cardinality > 1 and t.args[1].cardinality > 1:
            log.debug("can't do anything because we have multiple multivalued guys")
            return False
        if t.op == "If":
            log.debug("can't handle If")
            return False
        return True

    @staticmethod
    def _adjust_truism(t):
        """
        Swap the operands of the truism if the unknown variable is on the right side and the concrete value is on the
        left side.
        """
        if t.args[0].cardinality == 1 and t.args[1].cardinality > 1:
            return Balancer._reverse_comparison(t)
        return t

    #
    # Assumptions management
    #

    @staticmethod
    def _get_assumptions(t):
        """
        Given a constraint, _get_assumptions() returns a set of constraints that are implicitly
        assumed to be true. For example, `x <= 10` would return `x >= 0`.
        """

        if t.op in ("ULE", "ULT"):
            return [t.args[0] >= 0]
        if t.op in ("UGE", "UGT"):
            return [t.args[0] <= 2 ** len(t.args[0]) - 1]
        if t.op in ("SLE", "SLT"):
            return [claripy.SGE(t.args[0], -(1 << (len(t.args[0]) - 1)))]
        if t.op in ("SGE", "SGT"):
            return [claripy.SLE(t.args[0], (1 << (len(t.args[0]) - 1)) - 1)]
        return []

    #
    # Truism extractor
    #

    @staticmethod
    def _unpack_truisms(c) -> set:
        """
        Given a constraint, _unpack_truisms() returns a set of constraints that must be True for
        this constraint to be True.
        """
        match c.op:
            case "And":
                return Balancer._unpack_truisms_and(c)
            case "Not":
                return Balancer._unpack_truisms_not(c)
            case "Or":
                return Balancer._unpack_truisms_or(c)
            case _:
                return set()

    @staticmethod
    def _unpack_truisms_and(c):
        return set.union(*[Balancer._unpack_truisms(a) for a in c.args])

    @staticmethod
    def _unpack_truisms_not(c):
        if c.args[0].op == "And":
            return Balancer._unpack_truisms(claripy.Or(*[claripy.Not(a) for a in c.args[0].args]))
        if c.args[0].op == "Or":
            return Balancer._unpack_truisms(claripy.And(*[claripy.Not(a) for a in c.args[0].args]))
        return set()

    @staticmethod
    def _unpack_truisms_or(c):
        vals = [claripy.backends.vsa.is_false(v) for v in c.args]
        if all(vals):
            raise ClaripyBalancerUnsatError
        if vals.count(False) == 1:
            return Balancer._unpack_truisms(c.args[vals.index(False)])
        return set()

    #
    # Simplification routines
    #

    def _balance(self, truism):
        while True:
            log.debug("Balancing %s", truism)

            # can't balance single-arg bools (Not) for now
            if len(truism.args) == 1:
                return truism

            if not isinstance(truism.args[0], Base):
                return truism

            try:
                inner_aligned = Balancer._align_truism(truism)
                if inner_aligned.args[1].cardinality > 1:
                    log.debug("can't do anything because we have multiple multivalued guys")
                    return truism

                match inner_aligned.args[0].op:
                    case "Reverse":
                        balanced = Balancer._balance_reverse(inner_aligned)
                    case "__add__":
                        balanced = Balancer._balance_add(inner_aligned)
                    case "__sub__":
                        balanced = Balancer._balance_sub(inner_aligned)
                    case "ZeroExt":
                        balanced = Balancer._balance_zeroext(inner_aligned)
                    case "SignExt":
                        balanced = Balancer._balance_signext(inner_aligned)
                    case "Extract":
                        balanced = Balancer._balance_extract(inner_aligned)
                    case "__and__":
                        balanced = Balancer._balance_and(inner_aligned)
                    case "Concat":
                        balanced = Balancer._balance_concat(inner_aligned)
                    case "__lshift__":
                        balanced = Balancer._balance_lshift(inner_aligned)
                    case "If":
                        balanced = self._balance_if(inner_aligned)
                    case _:
                        log.debug("Balance handler %s not implemented.", truism.args[0].op)
                        return truism

                if balanced is inner_aligned:
                    return balanced
                truism = balanced
                continue
            except ClaripyBalancerError:
                log.warning("Balance handler for operation %s raised exception.", truism.args[0].op)
                return truism

    @staticmethod
    def _balance_reverse(truism):
        if truism.op in ["__eq__", "__ne__"]:
            return Bool(truism.op, (truism.args[0].args[0], truism.args[1].reversed))
        return truism

    @staticmethod
    def _balance_add(truism):
        if len(truism.args) != 2:
            return truism
        old_rhs = truism.args[1]
        lhs = truism.args[0]
        if all(a.concrete for a in lhs.args):
            # the old logic
            new_lhs = lhs.args[0]
            other_adds = lhs.args[1:]
        else:
            new_lhs = tuple(a for a in lhs.args if a.symbolic)
            if not new_lhs:
                return truism
            new_lhs = new_lhs[0] if len(new_lhs) == 1 else BV("__add__", new_lhs, length=len(lhs))
            other_adds = tuple(a for a in lhs.args if a.concrete)
            if not other_adds:
                return truism
        new_rhs = BV("__sub__", (old_rhs, *other_adds), length=len(lhs))
        return Bool(truism.op, (new_lhs, new_rhs))

    @staticmethod
    def _balance_sub(truism):
        if len(truism.args) != 2:
            return truism
        new_lhs = truism.args[0].args[0]
        old_rhs = truism.args[1]
        other_adds = truism.args[0].args[1:]
        new_rhs = BV("__add__", (old_rhs, *other_adds), length=len(truism.args[0]))
        return Bool(truism.op, (new_lhs, new_rhs))

    @staticmethod
    def _balance_zeroext(truism):
        num_zeroes, inner = truism.args[0].args
        other_side = truism.args[1][len(truism.args[1]) - 1 : len(truism.args[1]) - num_zeroes]

        if claripy.backends.vsa.is_true(other_side == 0):
            # We can safely eliminate this layer of ZeroExt
            return Bool(truism.op, (inner, truism.args[1][len(truism.args[1]) - num_zeroes - 1 : 0]))

        return truism

    @staticmethod
    def _balance_signext(truism):
        num_zeroes = truism.args[0].args[0]
        left_side = truism.args[0][len(truism.args[1]) - 1 : len(truism.args[1]) - num_zeroes]
        other_side = truism.args[1][len(truism.args[1]) - 1 : len(truism.args[1]) - num_zeroes]

        # TODO: what if this is a set value, but *not* the same as other_side
        if claripy.backends.vsa.identical(left_side, other_side):
            # We can safely eliminate this layer of ZeroExt
            return Bool(truism.op, (truism.args[0].args[1], truism.args[1][len(truism.args[1]) - num_zeroes - 1 : 0]))

        return truism

    @staticmethod
    def _balance_extract(truism):
        high, low, inner = truism.args[0].args
        inner_size = len(inner)

        if high < inner_size - 1:
            left_msb = inner[inner_size - 1 : high + 1]
            left_msb_zero = claripy.backends.vsa.is_true(left_msb == 0)
        else:
            left_msb = None
            left_msb_zero = None

        if low > 0:
            left_lsb = inner[high - 1 : 0]
            left_lsb_zero = claripy.backends.vsa.is_true(left_lsb == 0)
        else:
            left_lsb = None
            left_lsb_zero = None

        if left_msb_zero and left_lsb_zero:
            new_left = inner
            new_right = claripy.Concat(claripy.BVV(0, len(left_msb)), truism.args[1], claripy.BVV(0, len(left_lsb)))
            return Bool(truism.op, (new_left, new_right))
        if left_msb_zero:
            new_left = inner
            new_right = claripy.Concat(claripy.BVV(0, len(left_msb)), truism.args[1])
            return Bool(truism.op, (new_left, new_right))
        if left_lsb_zero:
            new_left = inner
            new_right = claripy.Concat(truism.args[1], claripy.BVV(0, len(left_lsb)))
            return Bool(truism.op, (new_left, new_right))

        if low == 0 and truism.args[1].op == "BVV" and truism.op not in {"SGE", "SLE", "SGT", "SLT"}:
            # single-valued rhs value with an unsigned operator
            # Eliminate Extract on lhs and zero-extend the value on rhs
            new_left = inner
            new_right = claripy.ZeroExt(inner.size() - truism.args[1].size(), truism.args[1])
            return Bool(truism.op, (new_left, new_right))

        return truism

    @staticmethod
    def _balance_and(truism):
        if len(truism.args[0].args) != 2:
            return truism
        op0, op1 = truism.args[0].args

        if op1.op == "BVV":
            # if all low bits of right are 1 and all high bits of right are 0, then this is equivalent to Extract()
            v = op1.args[0]
            low_ones = 0
            while v != 0:
                if v & 1 == 0:
                    # not all high bits are 0. abort
                    return truism
                low_ones += 1
                v >>= 1
            if low_ones == 0:
                # this should probably never happen
                new_left = BV("BVV", (0, truism.args[0].size()), length=truism.args[0].size())
                return Bool(truism.op, (new_left, truism.args[1]))

            if op0.op == "ZeroExt" and op0.args[0] + low_ones == op0.size():
                # ZeroExt(56, a) & 0xff == a  if a.size() == 8
                # we can safely remove __and__
                new_left = op0
                return Bool(truism.op, (new_left, truism.args[1]))

        return truism

    @staticmethod
    def _balance_concat(truism):
        size = len(truism.args[0])
        left_msb = truism.args[0].args[0]
        right_msb = truism.args[1][size - 1 : size - len(left_msb)]

        if claripy.backends.vsa.is_true(left_msb == 0) and claripy.backends.vsa.is_true(right_msb == 0):
            # we can cut these guys off!
            remaining_left = claripy.Concat(*truism.args[0].args[1:])
            remaining_right = truism.args[1][size - len(left_msb) - 1 : 0]
            return Bool(truism.op, (remaining_left, remaining_right))
        # TODO: handle non-zero single-valued cases
        return truism

    @staticmethod
    def _balance_lshift(truism):
        lhs = truism.args[0]
        rhs = truism.args[1]
        shift_amount_expr = lhs.args[1]
        expr = lhs.args[0]

        shift_amount_values = claripy.backends.vsa.eval(shift_amount_expr, 2)
        if len(shift_amount_values) != 1:
            return truism
        shift_amount = shift_amount_values[0]

        rhs_lower = claripy.Extract(shift_amount - 1, 0, rhs)
        rhs_lower_values = claripy.backends.vsa.eval(rhs_lower, 2)
        if len(rhs_lower_values) == 1 and rhs_lower_values[0] == 0:
            # we can remove the __lshift__

            return Bool(truism.op, (expr, rhs >> shift_amount))

        return truism

    def _balance_if(self, truism):
        condition, true_expr, false_expr = truism.args[0].args

        try:
            true_condition = getattr(true_expr, truism.op)(truism.args[1])
            false_condition = getattr(false_expr, truism.op)(truism.args[1])
        except ClaripyOperationError:
            # the condition was probably a Not (TODO)
            return truism

        can_true = claripy.backends.vsa.has_true(true_condition)
        can_false = claripy.backends.vsa.has_true(false_condition)
        must_true = claripy.backends.vsa.is_true(true_condition)
        must_false = claripy.backends.vsa.is_true(false_condition)

        if can_true and can_false:
            # always satisfiable
            return truism
        if not (can_true or can_false):
            # neither are satisfiable. This truism is fucked
            raise ClaripyBalancerUnsatError
        if must_true or (can_true and not can_false):
            # it will always be true
            self._truisms.append(condition)
            return Bool(truism.op, (true_expr, truism.args[1]))
        if must_false or (can_false and not can_true):
            # it will always be false
            self._truisms.append(~condition)
            return Bool(truism.op, (false_expr, truism.args[1]))
        return None

    #
    # Constraint handlers
    #

    def _handle(self, truism):
        log.debug("Handling %s", truism)

        if claripy.backends.vsa.is_false(truism):
            raise ClaripyBalancerUnsatError
        if Balancer._cardinality(truism.args[0]) == 1:
            # we are down to single-cardinality arguments, so our work is not
            # necessary
            return

        match truism.op:
            case "__eq__":
                self._handle_eq(truism)
            case "__ne__":
                self._handle_ne(truism)
            case "If":
                self._handle_if(truism)
            case "ULT" | "ULE" | "UGT" | "UGE" | "SLT" | "SLE" | "SGT" | "SGE":
                self._handle_comparison(truism)
            case _:
                log.debug("No handler for operation %s", truism.op)

    comparison_info = {  # noqa: RUF012
        "ULT": (True, False, True),
        "ULE": (True, True, True),
        "UGT": (False, False, True),
        "UGE": (False, True, True),
        "SLT": (True, False, False),
        "SLE": (True, True, False),
        "SGT": (False, False, False),
        "SGE": (False, True, False),
    }

    def _handle_comparison(self, truism):
        """
        Handles all comparisons.
        """

        is_lt, is_equal, is_unsigned = self.comparison_info[truism.op]

        size = len(truism.args[0])
        int_max = 2**size - 1 if is_unsigned else 2 ** (size - 1) - 1
        int_min = -(2 ** (size - 1))

        left_min = Balancer._min(truism.args[0], signed=not is_unsigned)
        left_max = Balancer._max(truism.args[0], signed=not is_unsigned)
        right_min = Balancer._min(truism.args[1], signed=not is_unsigned)
        right_max = Balancer._max(truism.args[1], signed=not is_unsigned)

        bound_max = right_max if is_equal else (right_max - 1 if is_lt else right_max + 1)
        bound_min = right_min if is_equal else (right_min - 1 if is_lt else right_min + 1)

        if is_lt and bound_max < int_min:
            # if the bound max is negative and we're unsigned less than, we're fucked
            raise ClaripyBalancerUnsatError
        if not is_lt and bound_min > int_max:
            # if the bound min is too big, we're fucked
            raise ClaripyBalancerUnsatError

        current_min = int_min
        current_max = int_max

        if is_lt:
            current_max = min(int_max, left_max, bound_max)
            self._add_upper_bound(truism.args[0], current_max)
        else:
            current_min = max(int_min, left_min, bound_min)
            self._add_lower_bound(truism.args[0], current_min)

    def _handle_eq(self, truism):
        lhs, rhs = truism.args
        if rhs.cardinality != 1:
            common = Balancer._same_bound_bv(lhs.intersection(rhs))
            mn, mx = Balancer._range(common)
            self._add_upper_bound(lhs, mx)
            self._add_upper_bound(rhs, mx)
            self._add_lower_bound(lhs, mn)
            self._add_lower_bound(rhs, mn)
        else:
            mn, mx = Balancer._range(rhs)
            self._add_upper_bound(lhs, mx)
            self._add_lower_bound(lhs, mn)

    def _handle_ne(self, truism):
        lhs, rhs = truism.args
        if rhs.cardinality == 1:
            val = claripy.backends.vsa.eval(rhs, 1)[0]
            max_int = vsa.StridedInterval.max_int(len(rhs))

            if val == 0:
                self._add_lower_bound(lhs, val + 1)
            elif val in (max_int, val - 1):
                self._add_upper_bound(lhs, max_int - 1)

    def _handle_if(self, truism):
        if claripy.backends.vsa.is_false(truism.args[2]):
            self._truisms.append(truism.args[0])
        elif claripy.backends.vsa.is_false(truism.args[1]):
            self._truisms.append(~truism.args[0])
