from .lexer import lex
from .parser import parse
from .exceptions import *
from .typeclass import TypeEnum, is_subtype_of
from .structs import *
from copy import deepcopy

def interpret(exprs, env, ns=None):
    ret = None
    for expr in exprs:
        ret = evaluate(expr, env, ns)
    return ret

def syntaxDef(symbol, arg, env, ns):
    res = evaluate(arg, env, ns)
    env.register(symbol.name, res)
    return List([])

def syntaxFunc(symbol, env, ns):
    func = UserFunction(symbol.name, symbol.args, symbol.many, symbol.body)
    func.return_type = symbol.return_type
    if symbol.name != "<lambda>":
        # register into the environment
        if env.contains(symbol.name):
            env_func = env.get(symbol.name)
            if isinstance(env_func, MultiFunction):
                env_func.register(func)
                env.register(symbol.name, env_func)
            else:
                raise InterpretPanic(symbol, "not a multifunc!")
        else:
            mf = MultiFunction(symbol.name)
            mf.register(func)
            env.register(symbol.name, mf)
        return List([])
    else:
        return func

def evaluate(expr, env, ns=None):
    if isinstance(expr, Literal) or isinstance(expr, Callable) or isinstance(expr, TypeWrap) or isinstance(expr, List) or isinstance(expr, Handle):
        return expr
    elif isinstance(expr, Symbol) or isinstance(expr, Type):
        if isinstance(expr, Symbol) and expr.quoted:
            return Symbol(expr.name, expr.line)
        if env.contains(expr.name):
            if isinstance(expr, Type) and expr.inner is not None:
                typecopy = deepcopy(env.get(expr.name))
                inner = env.get(f"{expr.inner}")
                typecopy.name.inner = inner
                return evaluate(typecopy, env, ns)
            else:
                return evaluate(env.get(expr.name), env, ns)
        elif ns is not None and env.contains(f"{ns}/{expr.name}"):
            if isinstance(expr, Type) and expr.inner is not None:
                typecopy = deepcopy(env.get(f"{ns}/{expr.name}"))
                inner = env.get(f"{expr.inner}")
                typecopy.name.inner = inner
                return evaluate(typecopy, env, ns)
            else:
                return evaluate(env.get(f"{ns}/{expr.name}"), env, ns)
        else:
            if isinstance(expr, Symbol):
                raise NebPanic(f"no such symbol: {expr}")
            else:
                raise NebPanic(f"no such type {expr}")

    # if it's an empty list, return it
    elif len(expr.args) == 0:
        return expr

    elif isinstance(expr.args[0], Expr):
        ev = evaluate(expr.args[0], env, ns)
        return evaluate(Expr([ev] + expr.args[1:]), env, ns)

    elif isinstance(expr.args[0], Callable):
        return expr.args[0].call(expr, env, ns)

    elif isinstance(expr.args[0], NebDef):
        return syntaxDef(expr.args[0], expr.args[1], env, ns)

    elif isinstance(expr.args[0], NebFuncDef):
        return syntaxFunc(expr.args[0], env, ns)

    elif not isinstance(expr.args[0], Symbol):
        raise NebPanic("can't evaluate without a symbol")
    name = expr.args[0].name
    if env.contains(name):
        return env.get(name).call(expr, env, ns)
    elif ns is not None and env.contains(f"{ns}/{name}"):
        return env.get(f"{ns}/{name}").call(expr, env, ns)
    else:
        raise InterpretPanic(expr.args[0], "unable to evaluate")


class Signature:

    def __init__(self, *args, many=None, return_type=None):
        self.args = args
        self.many = many
        if return_type is None:
            self.return_type = Type(":any")
    
    def compatable_arity(self, arity):
        if arity < len(self.args):
            return False
        elif self.many is None and arity > len(self.args):
            return False
        else:
            return True

    def arity_str(self):
        out = f"{len(self.args)}"
        if self.many is not None:
            out += "+"
        return out

    def get_type_by_idx(self, idx):
        if idx < len(self.args):
            return self.args[idx]
        else:
            return self.many

    def __str__(self):
        out = ""
        for arg in self.args:
            out = f"{out}{arg}"
        if self.many is not None:
            out = f"{out}{self.many}"
        return out


class Callable:

    def __init__(self, name, params, body, args=None, many=None):
        self.name = name
        self.params = params
        self.body = body
        self.args = args
        self.many = many
        self.type_ = Type(":any")  # TODO no it's not
        self.return_type = Type(":any")

        if args is None:
            self.sig = Signature(many=many)
        else:
            self.sig = Signature(*args, many=many)

    def describe(self, name=None):
        if name is None:
            name = self.name
        out = [f"({name}"]
        out.append(string_args(self.args, self.many))
        return " ".join(out).strip() + f") => {self.return_type}"

    def arity_check(self, symbol, params):
        min_arity = len([a for a in self.args if not a.optional])
        max_arity = -1 if self.many is not None else len(self.args)

        if len(params) < min_arity or (max_arity >= 0 and len(params) > max_arity):
            if max_arity < 0:
                fmt = f"{min_arity}+"
            elif min_arity != max_arity: 
                fmt = f"{min_arity}-{max_arity}"
            else:
                fmt = f"{min_arity}"
            raise InterpretPanic(symbol, f"expected [{fmt}] arguments, received {len(params)}")
        return True

    def call(self, expr, env):
        pass

class MultiFunction(Callable):

    def __init__(self, name):
        super().__init__(name, None, None)
        self.impls = {}

    def __str__(self):
        builtin_count = len(list(x for x in self.impls.values() if isinstance(x, Builtin)))
        userfunc_count = len(list(x for x in self.impls.values() if isinstance(x, UserFunction)))
        b = ""
        uf = ""
        if builtin_count != 0:
            b = f"{builtin_count} builtin"
        if userfunc_count != 0:
            uf = f"{userfunc_count} user defined"
        if b != "" and uf != "":
            desc = f"{b}, {uf}"
        else:
            desc = b if b != "" else uf
        return f"function {self.name} ({desc})"

    def describe(self):
        return [i.describe() for i in self.impls.values()]

    def user_impl(self):
        return len(list(x for x in self.impls.values() if isinstance(x, UserFunction))) != 0

    def register(self, impl):
        self.impls[f"{impl.sig}"] = impl

    def call(self, expr, env, ns):
        symbol = expr.args[0]
        params = expr.args[1:]

        # get compatable arities
        compatable_arities = [k for k,v in self.impls.items() if v.sig.compatable_arity(len(params))]
        if len(compatable_arities) == 0:
            fmt = "|".join(f"{v.sig.arity_str()}" for v in self.impls.values())
            raise InterpretPanic(symbol, f"expected [{fmt}] arguments, received {len(params)}")

        ret = []
        prev_candidates = []
        current_types = []
        next_candidates = compatable_arities
        for param_idx, param in enumerate(params):
            # evaluate the parameter
            ev = evaluate(param, env, ns)

            # reset the types we may be looking for
            current_types = []
            prev_candidates = next_candidates[:]
            next_candidates = []

            # loop through candidate functions
            for candidate in prev_candidates:
                func = self.impls[candidate]
                exp = func.sig.get_type_by_idx(param_idx)
                expected_type = evaluate(exp.type_, env, ns)
                current_types.append(expected_type)

                valid = expected_type.validate_type(ev, env, ns)
                if valid.value:
                    next_candidates.append(candidate)

            # if we have no more good functions, panic
            if len(next_candidates) == 0:
                fmt = "|".join(f"{t}" for t in current_types)
                rec = f"{ev.type_}"
                raise InterpretPanic(symbol, f"received {rec}, expected {fmt}", ev)
            else:
                ret.append(ev)

        if len(next_candidates) != 1:
            raise InterpretPanic(symbol, "ambiguous definition!")

        return self.impls[next_candidates[0]].call(Expr([symbol] + ret), env, ns)

class Special(Callable):

    def __init__(self, name, params, body, args=None, many=None):
        super().__init__(name, params, body, args, many)

class NebSyntax(Special):

    def __init__(self, name, callable_, args=None, many=None, return_type=None):
        super().__init__(name, None, callable_, args, many)
        if return_type is not None:
            self.return_type = return_type

    def __str__(self):
        return f"syntax function {self.name}"

    def call(self, expr, env, ns):
        return self.body(expr.args[0], expr.args[1:], env, ns)


class Builtin(Callable):
    
    def __init__(self, name, callable_, args=None, many=None, return_type=None):
        super().__init__(name, None, callable_, args, many)
        if return_type is not None:
            self.return_type = return_type

    def __str__(self):
        return f"builtin function {self.name}"

    def call(self, expr, env, ns):
        return self.body(expr.args[0], expr.args[1:], env, ns)


class UserFunction(Callable):
    
    def __init__(self, name, params, many, body):
        super().__init__(name, params, body, params, many)

    def __str__(self):
        out = f"(func {self.name} {self.return_type} ("
        args_list = [f"{a.name} {a.type_}" for a in self.args]
        if self.many:
            args_list.append(f"{self.many.name} {self.many.type_}")
        out = out + " ".join(args_list) + ") "
        for expr in self.body:
            out = out + f"{expr} "
        return out.strip() + ")"

    def call(self, expr, env, ns):
        evaluated_args = expr.args[1:]
        this_env = Environment(env)
        for idx, param in enumerate(self.params):
            this_env.register(param.name, evaluated_args[idx])

        # if we got "many", wrap the rest in a list
        if self.many:
            this_env.register(self.many.name, List(evaluated_args[len(self.params):]))
        return interpret(self.body, env=this_env, ns=ns)


class TypeWrap:

    def __init__(self, name, parent, is_func):
        self.name = ALL_TYPES[name]
        self.parent = parent
        self.is_func = is_func
        self.type_ = Type(":type")

    def validate_type(self, target, env, ns):
        # if it's an any type, it's valid
        if self.parent is None:
            return Bool(True)
        if isinstance(self.is_func, MultiFunction):
            valid = self.is_func.call(Expr([self.name, target]), env, ns)
        else:
            valid = self.is_func(self.name, [target], env, ns)
        if valid.value == True:
            return valid

        parent_type = env.get(target.type_.name)
        while valid.value != True and parent_type.parent is not None:
            parent_type = env.get(parent_type.parent.name.name)
            valid = Bool(self.name == parent_type.name) # TODO wrong
        return valid

    def __str__(self):
        return f"{self.name}"


class NebType(TypeWrap):
    pass

class UserType(TypeWrap):

    def __init__(self, name, parent, is_func):
        if name in ALL_TYPES:
            raise NebPanic(f"already a type called {name}")
        ALL_TYPES[name] = Type(name)
        super().__init__(name, parent, is_func)