#
# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
#

import ast
import builtins
import dis
import inspect
import pickle
import re
import sys
import textwrap
from collections import defaultdict, namedtuple
from logging import getLogger
from types import BuiltinFunctionType, CodeType, FunctionType, ModuleType
from typing import Any, Dict, List, Set, Tuple, Union

import opcode

# Python 3.8 needs to use typing.Iterable because collections.abc.Iterable is not subscriptable
# Python 3.9 can use both
# Python 3.10 needs to use collections.abc.Iterable because typing.Iterable is removed
if sys.version_info <= (3, 9):
    from typing import Iterable
else:
    from collections.abc import Iterable

logger = getLogger(__name__)

STORE_GLOBAL = opcode.opmap["STORE_GLOBAL"]
DELETE_GLOBAL = opcode.opmap["DELETE_GLOBAL"]
LOAD_GLOBAL = opcode.opmap["LOAD_GLOBAL"]
GLOBAL_OPS = (STORE_GLOBAL, DELETE_GLOBAL, LOAD_GLOBAL)

CODE_AS_COMMENT_HINT = (
    "The following comment contains the source code"
    " generated by snowpark-python for explanatory purposes.\n"
)
CODE_HEADER = """\
from __future__ import annotations
import pickle
"""
ImportNameAliasPair = namedtuple("ImportNameAliasPair", "name alias", defaults=[""] * 2)
ClassCodeGeneration = namedtuple(
    "ClassCodeGeneration", "class_object, generate_code", defaults=[None, True]
)


def get_func_references(func: FunctionType, ref_objects: Dict[str, Any]) -> None:
    """
    Get the objects references by target func, they could be methods, modules, classes, methods, global variables
    and its closures.
    This method will update the input ref_objects.

    Args:
        func: The target function to generate source code for.
        ref_objects: dict of objects referenced by the target function, key is the name and value is the object.
    """
    # 1. resolve function global references
    code_object = func.__code__
    globals_ref = extract_func_global_refs(
        func.__code__
    )  # get the names of the objects which func references
    globals = {
        k: func.__globals__[k]
        for k in globals_ref
        if k in func.__globals__  # retrieve the objects by names
        and k not in ref_objects
    }
    ref_objects.update(globals)

    # 2. resolve function closure references
    if func.__closure__ is not None:
        for k, v in zip(
            code_object.co_freevars,
            list(map(lambda x: x.cell_contents, func.__closure__)),
        ):
            ref_objects[k] = v
            # if the free variable is a function and is not itself (recursive) and is defined in the same module
            if (
                isinstance(v, FunctionType)
                and v != func
                and v.__module__ == func.__module__
            ):
                get_func_references(v, ref_objects)


def get_class_references(
    cls: type,
    func: FunctionType,
    ref_objects: Dict[str, Any],
    classes_to_generate: List[type],
    *,
    generate_code_for_class: bool = True,
) -> None:
    """
    To get the referenced objects of a class defined in the same module.
    A class could have methods, subclasses referencing other objects.
    This method will update the input ref_objects and classes_to_generate.

    Args:
        cls: The class to be analyzed to find references.
        func: The target function to generate source code for.
        ref_objects: dict of objects referenced by the target function, key is the name and value is the object.
        classes_to_generate: list of classes which are defined in the same module as the target function.
            Code generation is required for these classes.
        generate_code_for_class: Whether the source code for the class shall be generated.
    """
    func_module_name = func.__module__

    if generate_code_for_class:
        # order matters to classes_to_generate, when constructing source code,
        # referenced classes need to be defined first.
        classes_to_generate.insert(0, cls)

    # if base class is from the same module, we need to parse the class as well as generate code for the class
    # ClassCodeGeneration.generate_code indicates whether we should generate code for the class -- a nested class
    # doesn't need to be generated code for as the source code of the parent class contains it already.
    inferred_classes: List[ClassCodeGeneration] = []

    for base_class in cls.__bases__:
        base_class_name = base_class.__qualname__
        top_level_class = base_class_name.split(".")[0]
        if base_class_name == "object":
            continue
        elif base_class.__module__ == func_module_name:
            # this is a class that we should generate code for
            inferred_classes.append(
                ClassCodeGeneration(class_object=base_class, generate_code=True)
            )
        else:
            # if base class is from another module, we need to import it
            ref_objects[top_level_class] = func.__globals__[top_level_class]

    # __dict__ contains function, classmethod, classes attributes within a given class which
    # needs to be further analyzed
    for v in dict(cls.__dict__).values():
        if inspect.isclass(v):
            top_level_cls_name = v.__qualname__.split(".")[0]
            if v.__module__ == func_module_name:
                # v is a class defined in the same module as UDF func's, need to dynamically parse the class
                # one exception is that if v is class defined in cls,
                # then we should not re-generate code for the nested class
                inferred_classes.append(
                    ClassCodeGeneration(
                        class_object=v,
                        generate_code=not (
                            v.__qualname__.startswith(top_level_cls_name)
                        ),
                    )
                )
            else:
                # v is a class defined in another module, import the top level class from another module
                ref_objects[top_level_cls_name] = v
        elif inspect.isfunction(v) or isinstance(v, classmethod):
            # v is a function/classmethod, get the references objects of the Function object
            get_func_references(
                v if not isinstance(v, classmethod) else v.__func__, ref_objects
            )
        else:  # pragma: no cover
            # cls.__dict__ would also return __module__, __doc__, __weakref__ which are not required
            # for code generation, however, class variables is also included in __dict__, we don't do value evaluation
            # for them in the current implementation (e.g. the declaration of class variables is assigned the
            # result of function call). But we shall introduce an argument to control the behavior,
            # check JIRA SNOW-649884
            pass

    # recursively handling inferred classes that should be analyzed dynamically
    for inferred_class, generate_code in inferred_classes:
        # generate_code controls whether to generate code for inferred_class and is set to false
        # if inferred_classes is nested class.
        get_class_references(
            inferred_class,
            func,
            ref_objects,
            classes_to_generate,
            generate_code_for_class=generate_code,
        )


def extract_func_global_refs(code: CodeType) -> Set[str]:
    # inspired by cloudpickle to recursively extract all the global references used by the target func's code object
    # check: https://github.com/cloudpipe/cloudpickle/commit/6a0e12d058d1bd3ab26ec000ac2249b4ee7e9c9f
    out_names = set()
    for instr in dis.get_instructions(code):
        op = instr.opcode
        if op in GLOBAL_OPS:
            out_names.add(instr.argval)

    if code.co_consts:
        for const in code.co_consts:
            if isinstance(const, CodeType):
                out_names.update(extract_func_global_refs(const))

    return out_names


def remove_function_udf_annotation(udf_source_code: str) -> str:
    """
    Remove the udf/pandas_udf annotation to avoid re-registration.
    """
    udf_source_code = udf_source_code.strip()
    res = re.search(r"@(pandas_)?udf", udf_source_code)
    if res is None:
        return udf_source_code
    udf_anno_begin = res.start()
    udf_anno_end = res.end()
    if udf_source_code[udf_anno_end] == "\n":
        # just @udf
        return udf_source_code[udf_anno_end + 1 :]
    elif udf_source_code[udf_anno_end] != "(":
        # not a @udf
        return udf_source_code

    udf_anno_end = udf_anno_end + 1
    parenthesis_count = 1

    # find the pairing ')' of the leading 'udf('
    while parenthesis_count != 0:
        if udf_source_code[udf_anno_end] == "(":
            parenthesis_count += 1
        elif udf_source_code[udf_anno_end] == ")":
            parenthesis_count -= 1
        udf_anno_end += 1

    # check if there are still @udf annotations, raise error if there are still annotations
    code_after_remove = f"{udf_source_code[:udf_anno_begin].strip()}\n{udf_source_code[udf_anno_end:].strip()}".strip()
    if re.search(r"@(pandas_)?udf", code_after_remove) is not None:
        raise TypeError("An UDF can not be registered more than once.")
    return code_after_remove


def check_func_type(func: Any) -> None:
    """
    Check whether the target function is a valid type for source code generation. Raise error if not supported.
    """
    if (
        isinstance(func, classmethod)
        or inspect.ismethod(func)
        or not (isinstance(func, (FunctionType, BuiltinFunctionType)))
    ):
        error_msg = f"Code generation for {type(func)} is not supported yet."
        logger.debug(error_msg)
        raise TypeError(error_msg)


def generate_source_code(
    func: Union[FunctionType, BuiltinFunctionType], code_as_comment: bool = True
) -> str:
    """
    Dynamically generate source code of the given Python functions including:
      - The function itself
      - The functions/classes that are defined and referenced by the target function in the same module
      - The modules/class/method that have to be imported as used by the target function
      - The global or closure variables used by the target function

    The current implementation locks the support for the following scenarios:
     - Decorated functions: https://snowflakecomputing.atlassian.net/browse/SNOW-644983
     - Method and classmethod: https://snowflakecomputing.atlassian.net/browse/SNOW-644984

    Args:
        func: The target function to generate source code for.
        code_as_comment: Whether the code will be generated as comment.

    Returns:
        The generated source code.
    """

    try:
        check_func_type(func)
    except TypeError:
        if code_as_comment:
            # if it is an unsupported type, then no code generation and return empty string
            return ""
        raise

    # stored referenced object, key is the object name, value is the object
    ref_objects: Dict[str, Any] = {}
    # stored modules, each item should be a tuple of two strings, first is the true module name, second is the used name
    # such as alias or just the name
    to_import: Set[ImportNameAliasPair] = set()
    # imports class/funcs/vars from other modules, each key is the module name
    # each item is a set of tuples of two strings as the to_import, first module name, second alias
    to_import_from_module: Dict[str, Set[ImportNameAliasPair]] = defaultdict(set)
    # classes that should be generated in source code
    classes_to_generate: List[type] = []

    header_text = CODE_AS_COMMENT_HINT if code_as_comment else CODE_HEADER
    classes_text = ""

    # 1. find objects referenced by functions including classes, methods, modules, global variables
    find_target_func_objects_references(
        func, to_import, ref_objects, classes_to_generate
    )

    # 2. deal with the referenced objects by types
    func_text, global_vars_text = resolve_target_func_referenced_objects_by_type(
        func, to_import, to_import_from_module, ref_objects, code_as_comment
    )

    # 3. deal with the classes defined in the same module as func's
    for cls in classes_to_generate:
        classes_text = f"{classes_text}{textwrap.dedent(inspect.getsource(cls))}"

    # 4. deal with imports and alias
    imports_text = resolve_target_func_imports(to_import, to_import_from_module)

    # concatenating all the referenced parts
    source_code_without_target_func = f"{header_text}{imports_text}{global_vars_text}{classes_text}{func_text}".rstrip()

    # 5. handle func, remove the udf annotation
    complete_source_code, func_assignment = handle_target_func_self_source_code(
        func, source_code_without_target_func, code_as_comment
    )

    # 6. handle function assignment
    complete_source_code = f"""\
{complete_source_code}
{f"func = {func_assignment}"}\
""".strip()

    # 7. if code as comment is true, prefix each line with '#'
    if code_as_comment:
        complete_source_code = comment_source_code(complete_source_code)

    return complete_source_code.strip()


def is_lambda(func: FunctionType) -> bool:
    """
    Check whether the target function is a lambda function.
    """
    return func.__name__ == "<lambda>"


def get_lambda_code_text(code_text: str) -> str:
    """
    Extract the lambda expression from code text.

    Args:
        The original code text containing the lambda expression.

    Returns:
        The string of the lambda expression.

    """
    # add a wrapper to handle the case that the line of lambda source code does not include caller
    # such that ast could parse the expression tree:
    #     session.udf.register(
    #         lambda x, y: x + y, ...
    #     )
    try:
        source_ast = ast.parse(code_text)
    except SyntaxError as exc:
        if "cannot assign to lambda" in str(exc):
            # handle case like:
            # session.udf.register(
            #    lambda x, y: x + y, ...
            # )
            # code_text in this case is "lambda x, y: x + y, ..."
            code_text = f"wrapper({code_text})"
        # TODO: SNOW-685070 fix this
        elif "unmatched ')'" in str(exc):  # pragma: no cover
            # handle case like:
            # session.udf.register(
            #    lambda x, y: x + y, ...)
            # code_text in this case is "lambda x, y: x + y, ...)"
            code_text = f"wrapper({code_text}"
        source_ast = ast.parse(code_text)
    lambda_node = next(
        (node for node in ast.walk(source_ast) if isinstance(node, ast.Lambda)), None
    )
    if not lambda_node:
        raise TypeError("lambda function can not be extracted")

    lines = code_text.splitlines()
    # single line lambda
    if len(lines) == 1:
        return code_text[lambda_node.col_offset : lambda_node.end_col_offset]

    lambda_code_text = ""
    # lambda of multiple lines
    # handle case like:
    # session.udf.register(
    #    lambda x, y:\
    #    x + y, ...)
    for line_idx in range(lambda_node.lineno - 1, lambda_node.end_lineno):
        line = lines[line_idx]
        if line_idx == 0:
            lambda_code_text = f"{lambda_code_text}{line[lambda_node.col_offset:]}\n"
        elif line_idx == lambda_node.end_lineno - 1:
            lambda_code_text = f"{lambda_code_text}{line[: lambda_node.end_col_offset]}"
        else:
            lambda_code_text = f"{lambda_code_text}{line}\n"
    return lambda_code_text.strip()


def extract_submodule_imports(
    func: FunctionType, top_level_modules: Iterable[ModuleType]
) -> Set[ImportNameAliasPair]:
    """
    Get submodule imports, the func code co_names only gives the top level module names, the submodule imports
    have to be inferred manually.
    Top level modules refers the top level modules that is imported, e.g., in "import a1.a2",
    "a1" is the top level imported module.

    Consider the following example:

    import a1.a2.a3.a4
    def func():
        a1.a2.a3.a4.foo()

    func.__code__.co_names only contains ("a1", "a2", "a3", "a4", "foo") which does not include the
    complete import path information.

    To reconstruct "a1.a2.a3.a4", the current strategy is to import each prefix import of the import chains.
    This is not a perfect solution as we could import modules not used, but it works.

    Check https://snowflakecomputing.atlassian.net/browse/SNOW-651634 for more information.

    Args:
        func: The target function to generate source code for.
        top_level_modules: The name of top level modules from which to search the referenced imported objects
            or submodules.

    Returns:
        A set of tuple with each tuple composed of two string, the first one is actual name for the imported object,
        and the second one is alias.

    """
    imports = set()
    func_co_names = set(func.__code__.co_names)
    for module in top_level_modules:
        module_prefix = f"{module.__name__}."
        # search submodules that start with the top level module prefix
        # top level module is already collected by get_func_references method so will not be handled here
        for name in [m for m in sys.modules if m.startswith(module_prefix)]:
            tokens = set(name[len(module_prefix) :].split("."))
            # only add imports that co_names contains all the tokens
            # if any token does not show up in func_co_names, it means the module is not used by
            # func, thus there is no need to import the module
            # if all tokens showing up in the func_co_names, it means the module *might* be used by
            # func, import the token
            if not tokens - func_co_names:
                # submodule is not expected to have an alias.
                # alias module are expected to be detected and handled by get_func_references
                imports.add(ImportNameAliasPair(name=name))
    return imports


def find_target_func_objects_references(
    func: Union[FunctionType, BuiltinFunctionType],
    to_import: Set[ImportNameAliasPair],
    ref_objects: Dict[str, Any],
    classes_to_generate: List[type],
) -> None:
    """
    Find objects referenced by functions including classes, methods, modules, global variables.
    This method handles only FunctionType and BuiltinFunctionType. check_func_type will check the type first.

    Args:
        func: The target function to generate source code for.
        to_import: set of name and alias pairs of direct imports which should be generated as "import xxx"
            or "import xxx as yyy".
        ref_objects: dict of objects referenced by the target function, key is the name and value is the object.
        classes_to_generate: list of classes which are defined in the same module as the target function.
            Code generation is required for these classes.
    """
    func_module_name = func.__module__
    if isinstance(func, FunctionType):
        get_func_references(func, ref_objects)
        to_import.update(
            extract_submodule_imports(
                func, [v for v in ref_objects.values() if isinstance(v, ModuleType)]
            )
        )

        for v in ref_objects.values():
            if inspect.isclass(v) and v.__module__ == func_module_name:
                get_class_references(v, func, ref_objects, classes_to_generate)
    elif isinstance(func, BuiltinFunctionType):
        if func_module_name != builtins.__name__:
            to_import.add(ImportNameAliasPair(name=func_module_name))
    else:  # pragma: no cover
        raise TypeError(f"Code generation for {type(func)} is not supported yet.")


def resolve_target_func_referenced_objects_by_type(
    func: Union[FunctionType, BuiltinFunctionType],
    to_import: Set[ImportNameAliasPair],
    to_import_from_module: Dict[str, Set[ImportNameAliasPair]],
    ref_objects: Dict[str, Any],
    code_as_comment: bool,
) -> Tuple[str, str]:
    """
    Deal with the referenced objects by types, handles modules/classes/methods/global variables, generate source code
    for referenced functions defined in the same module as the target function's and referenced variables.

    Args:
        func: The target function to generate source code for.
        to_import: set of name and alias pairs of direct imports which should be generated as "import xxx"
            or "import xxx as yyy".
        to_import_from_module: dict of import information, key is the module name with value being the set of
            name and alias paris of imported objects which should be generated as "from xxx import yyy" or
            "from xxx import yyy as zzz".
        ref_objects: dict of objects referenced by the target function, key is the name and value is the object.
        code_as_comment: Whether the code will be generated as comment.

    Returns:
        A tuple of two strings, the first one is the source code of referenced functions defined in the same module
        as the target function's, and the second is the source code of referenced variables.
    """
    func_module_name = func.__module__
    func_text = ""
    global_vars_text = ""
    for name, obj in ref_objects.items():
        if obj == func:
            continue
        if inspect.ismodule(obj):
            # a) imported modules
            to_import.add(
                ImportNameAliasPair(
                    name=obj.__name__, alias=name if name != obj.__name__ else ""
                )
            )  # name could be an alias
        elif (
            inspect.isclass(obj) or inspect.isfunction(obj)
        ) and obj.__module__ != func_module_name:
            # b) classes or functions imported from other modules
            to_import_from_module[obj.__module__].add(
                ImportNameAliasPair(
                    name=obj.__name__, alias=name if name != obj.__name__ else ""
                )
            )  # name could be an alias
        else:
            # function/class/variables defined in the same module
            if inspect.isfunction(obj):
                func_text = f"{func_text}{textwrap.dedent(inspect.getsource(obj))}"
            elif inspect.isclass(obj):
                # dynamic class parsing will be handled separately
                continue
            else:
                # c) global variables used by UDF
                if code_as_comment:
                    global_vars_text = (
                        f"{global_vars_text}{name}  # variable of type {type(obj)}\n"
                    )
                    continue  # skip the serialization part if we just need code as comment
                # v does not have source code, then it's a global variable of which the value has been evaluated
                try:
                    global_vars_text = f"""\
{global_vars_text}
{name} = pickle.loads(bytes.fromhex('{pickle.dumps(obj).hex()}'))  \
# {name} is of type {type(obj)} and serialized by snowpark-python
"""
                except Exception as exc:  # pragma: no cover
                    logger.debug(
                        f"Unable to generate source code for object {name} of type {type(obj)} due to exception {exc}"
                    )
                    raise
    return func_text, global_vars_text


def resolve_target_func_imports(
    to_import: Set[ImportNameAliasPair],
    to_import_from_module: Dict[str, Set[ImportNameAliasPair]],
) -> str:
    """
    Deal with imports and alias, generate imports string.

    Args:
        to_import: set of name and alias pairs of direct imports which should be generated as "import xxx"
            or "import xxx as yyy".
        to_import_from_module: dict of import information, key is the module name with value being the set of
            name and alias paris of imported objects which should be generated as "from xxx import yyy" or
            "from xxx import yyy as zzz".

    Returns:
        A string of generated imports.
    """
    imports = [
        f"import {name + ' as ' if alias else ''}{alias or name}"
        for name, alias in sorted(to_import)
    ]
    for module, name_alias_pairs in sorted(to_import_from_module.items()):
        classes = ", ".join(
            f"{name + ' as ' if alias else ''}{alias or name}"
            for name, alias in sorted(name_alias_pairs)
        )
        imports.append(f"from {module} import {classes}")
    return "\n".join(imports) + ("\n" if imports else "")


def handle_target_func_self_source_code(
    func: Union[FunctionType, BuiltinFunctionType],
    source_code_without_target_func: str,
    code_as_comment: bool,
) -> Tuple[str, str]:
    """
    Generate the source code of the target func itself and apply function assignment.
    This method handles only FunctionType and BuiltinFunctionType. check_func_type will check the type first.

    Args:
        func: The target function to generate source code for.
        source_code_without_target_func: The generated code without the target function. The target function code
            and function assignment will be appended to this one.
        code_as_comment: Whether the code will be generated as comment.

    Returns:
        A tuple of two strings, the first one is the complete source code including target functions and all of its
        referenced objects, and the second one is function assignment.
    """
    func_module_name = func.__module__
    complete_source_code = source_code_without_target_func
    if isinstance(func, FunctionType):
        func_source_code = textwrap.dedent(inspect.getsource(func))
        if not code_as_comment:
            func_source_code = remove_function_udf_annotation(func_source_code)
        if not is_lambda(func):
            complete_source_code = f"{complete_source_code}\n{func_source_code}"
        func_assignment = (
            get_lambda_code_text(func_source_code) if is_lambda(func) else func.__name__
        )
    elif isinstance(func, BuiltinFunctionType):
        # BuiltinFunctionType
        func_assignment = (
            func.__name__
            if func_module_name == builtins.__name__
            else f"{func_module_name}.{func.__name__}"
        )
    else:  # pragma: no cover
        raise TypeError(f"Code generation for {type(func)} is not supported yet.")
    return complete_source_code, func_assignment


def comment_source_code(complete_source_code: str) -> str:
    """
    Prefix each line in source code with '#'

    Args:
        complete_source_code: The complete source code including target functions and all of its
        referenced objects

    Returns:
        The complete source code string with each line prefixed with "#".
    """
    return "\n".join(
        [f"#{f' {line}' if line else ''}" for line in complete_source_code.splitlines()]
    )
