from functools import wraps

import magma as m
from magma.bitutils import seq2int
from mantle import And, NAnd, Or, NOr, XOr, NXOr, LSL, LSR, Not, Invert, UDiv, SDiv, UMod, SMod, Mul
from mantle import ASR
from mantle import EQ, NE, ULT, ULE, UGT, UGE, SLT, SLE, SGT, SGE
from mantle import Mux
from .arith import Add, Sub, Negate

__all__ = []


def get_length(value):
    if isinstance(value, m._BitType):
        return None
    elif isinstance(value, m.ArrayType):
        return len(value)
    elif isinstance(value, int):
        return value.bit_length()
    else:
        raise NotImplementedError(f"Cannot get_length of {type(value)}")

def check_widths(width, args):
    for arg in args:
        arg_len = get_length(arg)
        # Allow integers with widths less than the width (they will be
        # promoted)
        if arg_len != width \
                and not ((isinstance(arg, int) and arg_len is 1 and width is None) or \
                         (isinstance(arg, int) and arg_len < width)):
            raise ValueError(
                f"All arguments should have the same length: {args}")


def check_operator_args(fn):
    @wraps(fn)
    def wrapped(*args, **kwargs):
        if len(args) < 2:
            raise RuntimeError(
                f"{fn.__name__} requires at least 2 arguments")
        width = get_length(args[0])
        check_widths(width, args)
        T = type(args[0])
        if not (all(issubclass(type(x).__class__, T.__class__) for x in args) or (
                all(issubclass(T.__class__, type(x).__class__) for x in args))):
            raise TypeError(
                "Currently Arguments to operators must be of the same type")
        return fn(width, *args, **kwargs)
    return wrapped


operators = {}


def _pass_closure_vars_as_args(*closure_args):
    def _wrapped(fn):
        @wraps(fn)
        def _wrapped_inner(*args, **kwargs):
            return fn(*closure_args, *args, **kwargs)
        return _wrapped_inner
    return _wrapped


def export(fn):
    global __all__
    __all__ += [fn.__name__]
    return fn


def preserve_type(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        retval = fn(*args, **kwargs)
        T = type(args[0])
        if isinstance(T, m.UIntKind):
            return m.uint(retval)
        elif isinstance(T, m.SIntKind):
            return m.sint(retval)
        return retval
    return wrapper


for _operator_name, _Circuit in (
    ("and_", And),
    ("nand", NAnd),
    ("or_", Or),
    ("nor", NOr),
    ("xor", XOr),
    ("nxor", NXOr),
    ("add", Add),
    ("sub", Sub),
    ("mul", Mul),
    ("udiv", UDiv),
    ("sdiv", SDiv),
    ("umod", UMod),
    ("smod", SMod)
):
    __all__ += [_operator_name]

    # Because Python uses dynamic scoping, the closures will use the
    # last value of _Circuit rather than capturing the lexical value.
    # Hacky workaround is to pass _Circuit as an argument to a
    # decorator so the "lexical" value is captured.
    @preserve_type
    @check_operator_args
    @_pass_closure_vars_as_args(_Circuit, _operator_name)
    def operator(circuit, name, width, *args, **kwargs):
        for arg in args:
            if not isinstance(arg, (m.Type, int)):
                raise TypeError("Operators only defined on magma types and Python ints")
        if name == "le" and args[0].isinput :
            assert len(args) == 2, "Expected two arguments for assignment (<=)"
            if not isinstance(args[1], int) and not args[1].isoutput():
                raise TypeError("Can only assign an output to an input")
        else:
            for arg in args:
                if not isinstance(arg, int) and not arg.isoutput():
                    raise TypeError("Non-assigment operators are only defined on output types")

        if name == "le" and args[0].isinput():
            # We could just call wire here, but we dispatch to the default
            # magma implementation so we only have to maintain the logic in one
            # place
            m.Type.__le__(args[0], args[1])
        else:
            if name in ["add", "sub", "mul", "udiv", "sdiv", 'umod', 'smod']:
                # These don't have a height
                if len(args) > 2:
                    raise Exception(f"{name} operator expects 2 arguments")
                return circuit(width, **kwargs)(*args)
            else:
                return circuit(len(args), width, **kwargs)(*args)
    operator.__name__ = _operator_name
    operator.__qualname__ = _operator_name
    operators[_operator_name] = operator
    exec(f"{_operator_name} = operator")


@export
@preserve_type
def lsl(I0, I1, **kwargs):
    width = get_length(I0)
    return LSL(width, **kwargs)(I0, I1)


@export
@preserve_type
def lsr(I0, I1, **kwargs):
    width = get_length(I0)
    return LSR(width, **kwargs)(I0, I1)


@export
@preserve_type
def asr(I0, I1, **kwargs):
    width = get_length(I0)
    return ASR(width, **kwargs)(I0, I1)


@export
@preserve_type
def not_(arg, **kwargs):
    return Not(**kwargs)(arg)


@export
@preserve_type
def invert(arg, **kwargs):
    width = get_length(arg)
    if width is None:
        return Not(**kwargs)(arg)
    else:
        return Invert(width, **kwargs)(arg)


@export
@preserve_type
def neg(arg, **kwargs):
    if isinstance(arg, int):
        return -arg
    return Negate(get_length(arg), **kwargs)(arg)


bitwise_ops = [
    ("__and__", and_),
    ("__or__", or_),
    ("__xor__", xor),
    ("__invert__", invert),
    ("__lshift__", lsl),
    ("__rshift__", lsr),
]

for method, op in bitwise_ops:
    if op not in (lsl, lsr):
        setattr(m.BitType, method, op)
    setattr(m.BitsType, method, op)

arithmetic_ops = [
    ("__neg__", neg),
    ("__add__", add),
    ("__sub__", sub),
    ("__mul__", mul),
]


@export
@check_operator_args
def lt(width, I0, I1, **kwargs):
    if isinstance(I0, m.SIntType):
        return SLT(width, **kwargs)(I0, I1)
    else:
        return ULT(width, **kwargs)(I0, I1)


@export
@check_operator_args
def le(width, I0, I1, **kwargs):
    if I0.isinput():
        return m.Type.__le__(I0, I1)
    if isinstance(I0, m.SIntType):
        return SLE(width, **kwargs)(I0, I1)
    else:
        return ULE(width, **kwargs)(I0, I1)


@export
@check_operator_args
def gt(width, I0, I1, **kwargs):
    if isinstance(I0, m.SIntType):
        return SGT(width, **kwargs)(I0, I1)
    else:
        return UGT(width, **kwargs)(I0, I1)


@export
@check_operator_args
def ge(width, I0, I1, **kwargs):
    if isinstance(I0, m.SIntType):
        return SGE(width, **kwargs)(I0, I1)
    else:
        return UGE(width, **kwargs)(I0, I1)


@export
@check_operator_args
def eq(width, I0, I1, **kwargs):
    return EQ(width, **kwargs)(I0, I1)


@export
@check_operator_args
def ne(width, I0, I1, **kwargs):
    return NE(width, **kwargs)(I0, I1)


relational_ops = [
    ("__lt__", lt),
    ("__le__", le),
    ("__gt__", gt),
    ("__ge__", ge),
]

for method, op in arithmetic_ops + relational_ops:
    setattr(m.SIntType, method, op)
    setattr(m.UIntType, method, op)

m.SIntType.__truediv__ = sdiv
m.UIntType.__truediv__ = udiv

m.SIntType.__mod__ = smod
m.UIntType.__mod__ = umod

for type_ in (m._BitType, m.ArrayType):
    setattr(type_, "__eq__", eq)
    setattr(type_, "__neq__", ne)


@export
@preserve_type
def mux(I, S, **kwargs):
    if isinstance(S, int):
        return I[S]
    elif S.const():
        return I[seq2int(S.bits())]
    return Mux(len(I), T=type(I[0]), **kwargs)(*I, S)


orig_get_item = m.ArrayType.__getitem__


def dynamic_mux_select(self, S):
    if isinstance(S, m.Type):
        if isinstance(self.T, m._BitKind):
            length = None
        else:
            length = len(self.T)
        height = len(self)
        inputs = self.ts[:]
        if m.mantle_target == "ice40" and height not in [2, 4, 8, 16]:
            if height > 16:
                raise NotImplementedError()
            orig_height = height
            height = 2 ** m.bitutils.clog2(height)
            if not isinstance(inputs[0], m._BitType):
                raise NotImplementedError(type(inputs[0]))
            inputs.extend([m.bit(0) for _ in range(height - orig_height)])

        return Mux(height, length)(*inputs, S)
    return orig_get_item(self, S)


setattr(m.ArrayType, "__getitem__", dynamic_mux_select)
