from __future__ import annotations
import logging

from archinfo.arch_soot import ArchSoot, SootAddressDescriptor, SootArgument, SootMethodDescriptor

from . import JNISimProcedure
from angr.calling_conventions import SimCCSoot
from angr.engines.soot.method_dispatcher import resolve_method

l = logging.getLogger("angr.procedures.java_jni.callmethod")

# pylint: disable=arguments-differ,unused-argument

#
# GetMethodID / GetStaticMethodID
#


class GetMethodID(JNISimProcedure):
    return_ty = "reference"

    def run(self, ptr_env, class_, ptr_method_name, ptr_method_sig):
        method_class = self.state.jni_references.lookup(class_)
        method_name = self._load_string_from_native_memory(ptr_method_name)
        method_sig = self._load_string_from_native_memory(ptr_method_sig)

        # derive parameter type from signature
        params, _ = ArchSoot.decode_method_signature(method_sig)

        # create method ID and return reference to it
        method_id = SootMethodDescriptor(method_class.name, method_name, params)
        return self.state.jni_references.create_new_reference(method_id)


#
# Call<Type>Method / CallNonvirtual<Type>Method / CallStatic<Type>Method
#


class CallMethodBase(JNISimProcedure):
    return_ty: str | None = None

    def _invoke(self, method_id, obj=None, dynamic_dispatch=True, args_in_array=None):
        # get invoke target
        class_name = obj.type if dynamic_dispatch else method_id.class_name
        invoke_target = resolve_method(self.state, method_id.name, class_name, method_id.params)
        invoke_addr = SootAddressDescriptor(invoke_target, 0, 0)

        # get args
        no_of_args = len(invoke_target.params)
        if args_in_array is not None:
            arg_values = self._get_arg_values_from_array(args_in_array, no_of_args)
        else:
            arg_values = self._get_arg_values(no_of_args)

        # setup java args
        java_args = self._setup_java_args(arg_values, invoke_target, this_ref=obj)

        # call java method
        # => after returning, the execution will be continued in _return_result_of_computation
        self.call(invoke_addr, java_args, "return_from_invocation", cc=SimCCSoot(ArchSoot()))

    def _get_arg_values(self, no_of_args):
        return [self.va_arg("void*") for _ in range(no_of_args)]

    def _get_arg_values_from_array(self, array, no_of_args):
        return self._load_from_native_memory(
            addr=array, data_size=self.arch.bytes, no_of_elements=no_of_args, return_as_list=True
        )

    def _setup_java_args(self, arg_values, method_id, this_ref=None):
        args = []

        # if available, add 'this' reference
        if this_ref:
            args += [SootArgument(this_ref, this_ref.type, is_this_ref=True)]

        # function arguments
        for arg_value_, arg_type in zip(arg_values, method_id.params):
            if arg_type in ArchSoot.primitive_types:
                # argument has a primitive integral type
                # => cast native value to java type
                arg_value = self.project.simos.cast_primitive(self.state, value=arg_value_, to_type=arg_type)

            else:
                # argument has a relative type
                # => lookup java object
                arg_value = self.state.jni_references.lookup(arg_value_)

            args += [SootArgument(arg_value, arg_type)]

        return args

    def _return_from_invocation(self):
        if self.return_ty != "void":
            ret_value = self.state.javavm_registers.load("invoke_return_value")
            if self.return_ty == "reference":
                return self.state.jni_references.create_new_reference(ret_value)
            return ret_value
        return None


#
# Call<Type>Method
#


class CallMethod(CallMethodBase):
    def run(self, ptr_env, obj_, method_id_):
        method_id = self.state.jni_references.lookup(method_id_)
        obj = self.state.jni_references.lookup(obj_)
        self._invoke(method_id, obj, dynamic_dispatch=True)

    def return_from_invocation(self, ptr_env, obj_, method_id_):
        return self._return_from_invocation()


class CallMethodA(CallMethodBase):
    def run(self, ptr_env, obj_, method_id_, ptr_args):
        method_id = self.state.jni_references.lookup(method_id_)
        obj = self.state.jni_references.lookup(obj_)
        self._invoke(method_id, obj, dynamic_dispatch=True, args_in_array=ptr_args)

    def return_from_invocation(self, ptr_env, obj_, method_id_, ptr_args):
        return self._return_from_invocation()


class CallObjectMethod(CallMethod):
    return_ty = "reference"


class CallBooleanMethod(CallMethod):
    return_ty = "boolean"


class CallByteMethod(CallMethod):
    return_ty = "byte"


class CallCharMethod(CallMethod):
    return_ty = "char"


class CallShortMethod(CallMethod):
    return_ty = "short"


class CallIntMethod(CallMethod):
    return_ty = "int"


class CallLongMethod(CallMethod):
    return_ty = "long"


class CallVoidMethod(CallMethod):
    return_ty = "void"


class CallObjectMethodA(CallMethodA):
    return_ty = "reference"


class CallBooleanMethodA(CallMethodA):
    return_ty = "boolean"


class CallByteMethodA(CallMethodA):
    return_ty = "byte"


class CallCharMethodA(CallMethodA):
    return_ty = "char"


class CallShortMethodA(CallMethodA):
    return_ty = "short"


class CallIntMethodA(CallMethodA):
    return_ty = "int"


class CallLongMethodA(CallMethodA):
    return_ty = "long"


class CallVoidMethodA(CallMethodA):
    return_ty = "void"


#
# CallNonVirtual<Type>Method
#


class CallNonvirtualMethod(CallMethodBase):
    def run(self, ptr_env, obj_, class_, method_id_):
        method_id = self.state.jni_references.lookup(method_id_)
        obj = self.state.jni_references.lookup(obj_)
        self._invoke(method_id, obj, dynamic_dispatch=False)

    def return_from_invocation(self, ptr_env, obj_, class_, method_id_):
        return self._return_from_invocation()


class CallNonvirtualMethodA(CallMethodBase):
    def run(self, ptr_env, obj_, method_id_, ptr_args):
        method_id = self.state.jni_references.lookup(method_id_)
        obj = self.state.jni_references.lookup(obj_)
        self._invoke(method_id, obj, dynamic_dispatch=False, args_in_array=ptr_args)

    def return_from_invocation(self, ptr_env, obj_, method_id_, ptr_args):
        return self._return_from_invocation()


class CallNonvirtualObjectMethod(CallNonvirtualMethod):
    return_ty = "reference"


class CallNonvirtualBooleanMethod(CallNonvirtualMethod):
    return_ty = "boolean"


class CallNonvirtualByteMethod(CallNonvirtualMethod):
    return_ty = "byte"


class CallNonvirtualCharMethod(CallNonvirtualMethod):
    return_ty = "char"


class CallNonvirtualShortMethod(CallNonvirtualMethod):
    return_ty = "short"


class CallNonvirtualIntMethod(CallNonvirtualMethod):
    return_ty = "int"


class CallNonvirtualLongMethod(CallNonvirtualMethod):
    return_ty = "long"


class CallNonvirtualVoidMethod(CallNonvirtualMethod):
    return_ty = "void"


class CallNonvirtualObjectMethodA(CallNonvirtualMethodA):
    return_ty = "reference"


class CallNonvirtualBooleanMethodA(CallNonvirtualMethodA):
    return_ty = "boolean"


class CallNonvirtualByteMethodA(CallNonvirtualMethodA):
    return_ty = "byte"


class CallNonvirtualCharMethodA(CallNonvirtualMethodA):
    return_ty = "char"


class CallNonvirtualShortMethodA(CallNonvirtualMethodA):
    return_ty = "short"


class CallNonvirtualIntMethodA(CallNonvirtualMethodA):
    return_ty = "int"


class CallNonvirtualLongMethodA(CallNonvirtualMethodA):
    return_ty = "long"


class CallNonvirtualVoidMethodA(CallNonvirtualMethodA):
    return_ty = "void"


#
# CallStatic<Type>Method
#


class CallStaticMethod(CallMethodBase):
    def run(self, ptr_env, class_, method_id_):
        method_id = self.state.jni_references.lookup(method_id_)
        self._invoke(method_id, dynamic_dispatch=False)

    def return_from_invocation(self, ptr_env, class_, method_id_):
        return self._return_from_invocation()


class CallStaticMethodA(CallMethodBase):
    def run(self, ptr_env, obj_, method_id_, ptr_args):
        method_id = self.state.jni_references.lookup(method_id_)
        self._invoke(method_id, dynamic_dispatch=False, args_in_array=ptr_args)

    def return_from_invocation(self, ptr_env, class_, method_id_, ptr_args):
        return self._return_from_invocation()


class CallStaticObjectMethod(CallStaticMethod):
    return_ty = "reference"


class CallStaticBooleanMethod(CallStaticMethod):
    return_ty = "boolean"


class CallStaticByteMethod(CallStaticMethod):
    return_ty = "byte"


class CallStaticCharMethod(CallStaticMethod):
    return_ty = "char"


class CallStaticShortMethod(CallStaticMethod):
    return_ty = "short"


class CallStaticIntMethod(CallStaticMethod):
    return_ty = "int"


class CallStaticLongMethod(CallStaticMethod):
    return_ty = "long"


class CallStaticVoidMethod(CallStaticMethod):
    return_ty = "void"


class CallStaticObjectMethodA(CallStaticMethodA):
    return_ty = "reference"


class CallStaticBooleanMethodA(CallStaticMethodA):
    return_ty = "boolean"


class CallStaticByteMethodA(CallStaticMethodA):
    return_ty = "byte"


class CallStaticCharMethodA(CallStaticMethodA):
    return_ty = "char"


class CallStaticShortMethodA(CallStaticMethodA):
    return_ty = "short"


class CallStaticIntMethodA(CallStaticMethodA):
    return_ty = "int"


class CallStaticLongMethodA(CallStaticMethodA):
    return_ty = "long"


class CallStaticVoidMethodA(CallStaticMethodA):
    return_ty = "void"
