from __future__ import annotations
import itertools
from angr.ailment.expression import BinaryOp, Const, Convert, Extract, Insert
from .base import PeepholeOptimizationExprBase


class SimplifyBitwiseInserts(PeepholeOptimizationExprBase):
    """
    simplifies _INSERT(A, c, EXTRACT(A, ...) | B) => A | B
    """

    __slots__ = ()

    NAME = "Simplify bitwise inserts"
    expr_classes = (Insert,)

    def optimize(self, expr, **kwargs):
        pb1l = [(0, expr.base)]
        if isinstance(expr.base, BinaryOp) and expr.base.op == "Concat":
            if isinstance(expr.base.operands[0], Const) and expr.base.operands[0].value == 0:
                pb1l.append((expr.base.bits, expr.base.operands[1]))
            else:
                pb1l.append((0, expr.base.operands[0]))
        elif isinstance(expr.base, Convert) and expr.base.operand.bits < expr.base.bits:
            # I don't know if 0 is right here
            pb1l.append((0, expr.base.operand))

        if not isinstance(expr.value, BinaryOp):
            return None
        if expr.value.op not in ("Or", "Xor"):
            return None
        pb2l = []
        for potential_base, other_base in zip(expr.value.operands, reversed(expr.value.operands)):
            pb2l.append((0, potential_base, other_base))
            if isinstance(other_base, Const):
                if isinstance(potential_base, Extract) and isinstance(potential_base.offset, Const):
                    pb2l.append((potential_base.offset, potential_base.base, other_base))
                elif isinstance(potential_base, Convert):
                    # I don't know if 0 is right here
                    pb2l.append((0, potential_base.operand, other_base))

        for (pb1o, pb1), (pb2o, pb2, pb2x) in itertools.product(pb1l, pb2l):  # noqa: B007
            if pb1o == pb2o and pb1.bits == pb2.bits and pb1.likes(pb2):
                break
        else:
            return None

        if pb1o != 0:
            # hello future angr developer! This branch exists because I don't have an example to exercise this branch.
            # for offsets != 0. If you see code enter this branch, please remove this and test for correctness :)
            return None

        if pb2x.bits != pb2.bits:
            # should only be generated by guarded logic above
            assert pb2.bits > pb2x.bits
            assert isinstance(pb2x, Const)
            assert isinstance(pb2x.value, int)
            pb2x = Const(self.manager.next_atom(), None, pb2x.value << pb1o, pb2.bits)
        elif pb1o != 0:
            pb2x = BinaryOp(
                self.manager.next_atom(), "Shl", [pb2x, Const(self.manager.next_atom(), None, pb1o, pb2x.bits)]
            )

        result = BinaryOp(self.manager.next_atom(), expr.value.op, [pb1, pb2x])
        if result.bits != expr.bits:
            result = Convert(self.manager.next_atom(), result.bits, expr.bits, False, result)
        return result
