
import functools

from .vex_helper import IRSBCustomizer, Type, JumpKind
import pyvex
from pyvex.expr import IRExpr, Const, RdTmp, Unop, Binop, Triop, Qop, Load, CCall, Get


def checkparams(*fn, **options):
    rhstype = options.pop('rhstype', None)
    if options:
        raise TypeError("unsupported keyword arguments: %s" %
                        ",".join(options.keys()))

    if fn:
        def wrapper(self, *args, **kwargs):
            irsb_cs = {a.irsb_c for a in list(args) + [self] if isinstance(a, VexValue)}  # pylint: disable=no-member
            assert len(irsb_cs) == 1, 'All VexValues must belong to the same irsb_c'
            args = list(args)
            for arg in args:
                if isinstance(arg, int):
                    thetype = rhstype if rhstype else self.ty
                    args[args.index(arg)] = VexValue.Constant(self.irsb_c, arg, thetype)
            args = tuple(args)
            return fn[0](self, *args, **kwargs)
        return functools.update_wrapper(wrapper, fn[0])
    else:
        def decorator(fn):
            def wrapper(self, *args, **kwargs):
                irsb_cs = {a.irsb_c for a in list(args) + [self] if
                           isinstance(a, VexValue)}  # pylint: disable=no-member
                assert len(irsb_cs) == 1, 'All VexValues must belong to the same irsb_c'
                args = list(args)
                for arg in args:
                    if isinstance(arg, int):
                        thetype = rhstype if rhstype else self.ty
                        args[args.index(arg)] = VexValue.Constant(self.irsb_c, arg, thetype)
                args = tuple(args)
                return fn(self, *args, **kwargs)
            return functools.update_wrapper(wrapper, fn)
        return decorator


def vvifyresults(f):
    @functools.wraps(f)
    def decor(self, *args, **kwargs):
        returned = f(self, *args, **kwargs)
        assert isinstance(returned, RdTmp) or isinstance(returned, Const)
        return VexValue(self.irsb_c, returned)
    return decor


class VexValue:
    def __init__(self, irsb_c, rdt):
        self.irsb_c = irsb_c
        self.ty = self.irsb_c.get_type(rdt)
        self.rdt = rdt
        self.width = pyvex.get_type_size(self.ty)
        self.is_signed = False

    @property
    def value(self):
        if isinstance(self.rdt, Const):
            return self.rdt.con.value
        else:
            raise ValueError("Non-constant VexValue has no value property")

    def toSigned(self):
        self.is_signed = True

    def toUnsigned(self):
        self.is_signed = False

    @vvifyresults
    def widen_unsigned(self, ty):
        return self.irsb_c.op_widen_int_unsigned(self.rdt, ty)

    @vvifyresults
    def cast_to(self, ty, signed=False, high=False):
        return self.irsb_c.cast_to(self.rdt, ty, signed=signed, high=high)

    @vvifyresults
    def widen_signed(self, ty):
        return self.irsb_c.op_widen_int_signed(self.rdt, ty)

    @vvifyresults
    def narrow_high(self, ty):
        return self.irsb_c.op_narrow_int(self.rdt, ty, high_half=True)

    @vvifyresults
    def narrow_low(self, ty):
        return self.irsb_c.op_narrow_int(self.rdt, ty, high_half=False)

    # TODO at some point extend this to Vex nonconstants
    def __getitem__(self, idx):
        getb = lambda i: VexValue(self.irsb_c, self.irsb_c.get_bit(self.rdt, i))
        makeconstant = lambda x: VexValue.Constant(self.irsb_c, x, Type.int_8).rdt
        if not isinstance(idx, slice):
            actualindex = slice(idx).indices(self.width)[1]
            return getb(makeconstant(actualindex))
        else:
            return [getb(makeconstant(i)) for i in xrange(*idx.indices(self.width))]

    @checkparams
    @vvifyresults
    def set_bit(self, idx, bval):
        typedidx = idx.cast_to(Type.int_8)
        return self.irsb_c.set_bit(self.rdt, idx.rdt, bval.rdt)

    @checkparams
    @vvifyresults
    def set_bits(self, idxsandvals):
        return self.irsb_c.set_bits(self.rdt, [(i.cast_to(Type.int_8).rdt, b.rdt) for i, b in idxsandvals])

    @checkparams
    @vvifyresults
    def ite(self, iftrue, iffalse):
        onebitcond = self.cast_to(Type.int_1)
        return self.irsb_c.ite(onebitcond.rdt, iftrue.rdt, iffalse.rdt)

    @checkparams
    @vvifyresults
    def __add__(self, right):
        return self.irsb_c.op_add(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __sub__(self, right): # I'm not sure if we want rsub
        return self.irsb_c.op_sub(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __div__(self, right):
        return self.irsb_c.op_div(self.rdt, right.rdt, self.is_signed)

    @checkparams
    @vvifyresults
    def __floordiv__(self, right): # Note: nonprimitive
        return self / right

    @checkparams
    @vvifyresults
    def __truediv__(self, right): # Note: nonprimitive
        return self / right

    @checkparams
    @vvifyresults
    def __and__(self, right):
        return self.irsb_c.op_and(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __eq__(self, right):
        return self.irsb_c.op_cmp_eq(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __ne__(self, other):
        return self.irsb_c.op_cmp_ne(self.rdt, other.rdt)

    @checkparams
    @vvifyresults
    def __invert__(self):
        return self.irsb_c.op_not(self.rdt)

    @checkparams
    @vvifyresults
    def __le__(self, right):
        if self.is_signed:
            return self.irsb_c.op_cmp_sle(self.rdt, right.rdt)
        else:
            return self.irsb_c.op_cmp_ule(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __gt__(self, other):
        if self.is_signed:
            return self.irsb_c.op_cmp_sgt(self.rdt, other.rdt)
        else:
            return self.irsb_c.op_cmp_ugt(self.rdt, other.rdt)

    @checkparams
    @vvifyresults
    def __ge__(self, right):
        if self.is_signed:
            return self.irsb_c.op_cmp_sge(self.rdt, right.rdt)
        else:
            return self.irsb_c.op_cmp_uge(self.rdt, right.rdt)

    @checkparams(rhstype=Type.int_8)
    @vvifyresults
    def __lshift__(self, right):
        return self.irsb_c.op_shl(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __lt__(self, right):
        if self.is_signed:
            return self.irsb_c.op_cmp_slt(self.rdt, right.rdt)
        else:
            return self.irsb_c.op_cmp_ult(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __mod__(self, right): # Note: nonprimitive
        return self.rdt - (self.rdt / right.rdt) * right.rdt

    @checkparams
    @vvifyresults
    def __mul__(self, right):
        return self.irsb_c.op_mull(self.rdt, right.rdt, signed=self.is_signed)

    @checkparams
    @vvifyresults
    def __neg__(self): # Note: nonprimitive
        if self.is_unsigned:
            raise Exception('Number is unsigned, cannot change sign!')
        else:
            return self.rdt * -1

    @checkparams
    @vvifyresults
    def __or__(self, right):
        return self.irsb_c.op_or(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __pos__(self):
        return self

    @checkparams
    @vvifyresults
    def __rshift__(self, right):
        return self.irsb_c.op_shr(self.rdt, right.rdt)

    @checkparams
    @vvifyresults
    def __xor__ (self, right):
        return self.irsb_c.op_xor(self.rdt, right.rdt)

    @classmethod
    def Constant(cls, irsb_c, val, ty):
        """
        Creates a constant as a VexValue
        :param irsb_c: The IRSBCustomizer to use
        :param val: The value, as an integer
        :param ty: The type of the resulting VexValue
        :return: a VexValue
        """
        assert not (isinstance(val, VexValue) or isinstance(val, IRExpr))
        rdt = irsb_c.mkconst(val, ty)
        return cls(irsb_c, rdt)

