import functools

from vyper import (
    ast as vy_ast,
)
from vyper.exceptions import (
    InvalidLiteral,
    StructureException,
    TypeMismatch,
)
from vyper.parser.expr import (
    Expr,
)
from vyper.types import (
    BaseType,
    ByteArrayType,
    StringType,
    is_base_type,
)
from vyper.utils import (
    SizeLimits,
)


class Optional(object):
    def __init__(self, typ, default):
        self.typ = typ
        self.default = default


def process_arg(index, arg, expected_arg_typelist, function_name, context):
    if isinstance(expected_arg_typelist, Optional):
        expected_arg_typelist = expected_arg_typelist.typ
    if not isinstance(expected_arg_typelist, tuple):
        expected_arg_typelist = (expected_arg_typelist, )

    vsub = None
    for expected_arg in expected_arg_typelist:
        if expected_arg == 'num_literal':
            if context.constants.is_constant_of_base_type(arg, ('uint256', 'int128')):
                return context.constants.get_constant(arg.id, None).value
            if isinstance(arg, (vy_ast.Int, vy_ast.Decimal)):
                return arg.n
        elif expected_arg == 'str_literal':
            if isinstance(arg, vy_ast.Str):
                bytez = b''
                for c in arg.s:
                    if ord(c) >= 256:
                        raise InvalidLiteral(
                            f"Cannot insert special character {c} into byte array",
                            arg,
                        )
                    bytez += bytes([ord(c)])
                return bytez
        elif expected_arg == 'bytes_literal':
            if isinstance(arg, vy_ast.Bytes):
                return arg.s
        elif expected_arg == 'name_literal':
            if isinstance(arg, vy_ast.Name):
                return arg.id
            elif isinstance(arg, vy_ast.Subscript) and arg.value.id == 'bytes':
                return f'bytes[{arg.slice.value.n}]'
        elif expected_arg == '*':
            return arg
        elif expected_arg == 'bytes':
            sub = Expr(arg, context).lll_node
            if isinstance(sub.typ, ByteArrayType):
                return sub
        elif expected_arg == 'string':
            sub = Expr(arg, context).lll_node
            if isinstance(sub.typ, StringType):
                return sub
        else:
            # Does not work for unit-endowed types inside compound types, e.g. timestamp[2]
            parsed_expected_type = context.parse_type(
                vy_ast.parse_to_ast(expected_arg)[0].value,
                'memory',
            )
            if isinstance(parsed_expected_type, BaseType):
                vsub = vsub or Expr.parse_value_expr(arg, context)

                is_valid_integer = (
                    (
                        expected_arg in ('int128', 'uint256') and isinstance(vsub.typ, BaseType)
                    ) and (
                        vsub.typ.typ in ('int128', 'uint256') and vsub.typ.is_literal
                    ) and (
                        SizeLimits.in_bounds(expected_arg, vsub.value)
                    )
                )

                if is_base_type(vsub.typ, expected_arg):
                    return vsub
                elif is_valid_integer:
                    return vsub
            else:
                vsub = vsub or Expr(arg, context).lll_node
                if vsub.typ == parsed_expected_type:
                    return Expr(arg, context).lll_node
    if len(expected_arg_typelist) == 1:
        raise TypeMismatch(
            f"Expecting {expected_arg} for argument {index} of {function_name}",
            arg
        )
    else:
        raise TypeMismatch(
            f"Expecting one of {expected_arg_typelist} for argument {index} of {function_name}",
            arg
        )


def signature(*argz, **kwargz):
    def decorator(f):
        @functools.wraps(f)
        def g(element, context):
            function_name = element.func.id
            if len(element.args) > len(argz):
                raise StructureException(
                    f"Expected {len(argz)} arguments for {function_name}, "
                    f"got {len(element.args)}",
                    element
                )
            subs = []
            for i, expected_arg in enumerate(argz):
                if len(element.args) > i:
                    subs.append(process_arg(
                        i + 1,
                        element.args[i],
                        expected_arg,
                        function_name,
                        context,
                    ))
                elif isinstance(expected_arg, Optional):
                    subs.append(expected_arg.default)
                else:
                    raise StructureException(
                        f"Not enough arguments for function: {element.func.id}",
                        element
                    )
            kwsubs = {}
            element_kw = {k.arg: k.value for k in element.keywords}
            for k, expected_arg in kwargz.items():
                if k not in element_kw:
                    if isinstance(expected_arg, Optional):
                        kwsubs[k] = expected_arg.default
                    else:
                        raise StructureException(f"Function {function_name} requires argument {k}",
                                                 element)
                else:
                    kwsubs[k] = process_arg(k, element_kw[k], expected_arg, function_name, context)
            for k, _arg in element_kw.items():
                if k not in kwargz:
                    raise StructureException(f"Unexpected argument: {k}",
                                             element)
            return f(element, subs, kwsubs, context)
        return g
    return decorator
