"""

AutoProfile Script Demo
=======================

The following demo is end-to-end bash code that writes a demo script and
profiles it with autoprofile.

.. code:: bash

    # Write demo python script to disk
    python -c "if 1:
        import textwrap
        text = textwrap.dedent(
            '''
            def plus(a, b):
                return a + b

            def fib(n):
                a, b = 0, 1
                while a < n:
                    a, b = b, plus(a, b)

            def main():
                import math
                import time
                start = time.time()

                print('start calculating')
                while time.time() - start < 1:
                    fib(10)
                    math.factorial(1000)
                print('done calculating')

            main()
            '''
        ).strip()
        with open('demo.py', 'w') as file:
            file.write(text)
    "

    echo "---"
    echo "## Profile With AutoProfile"
    python -m kernprof -p demo.py -l demo.py
    python -m line_profiler -rmt demo.py.lprof
"""

from __future__ import annotations
import contextlib
import functools
import importlib.util
import operator
import sys
import types
from collections.abc import MutableMapping
from typing import Any, cast, Dict, Mapping
from typing import ContextManager
from .ast_tree_profiler import AstTreeProfiler
from .run_module import AstTreeModuleProfiler
from .line_profiler_utils import add_imported_function_or_module
from .util_static import modpath_to_modname

PROFILER_LOCALS_NAME = 'prof'


def _extend_line_profiler_for_profiling_imports(prof: Any) -> None:
    """Allow profiler to handle functions/methods, classes & modules with a single call.

    Add a method to LineProfiler that can identify whether the object is a
    function/method, class or module and handle it's profiling accordingly.
    Mainly used for profiling objects that are imported.
    (Workaround to keep changes needed by autoprofile separate from base LineProfiler)

    Args:
        prof (LineProfiler):
            instance of LineProfiler.
    """
    prof.add_imported_function_or_module = types.MethodType(
        add_imported_function_or_module, prof
    )


def run(
    script_file: str,
    ns: MutableMapping[str, Any],
    prof_mod: list[str],
    profile_imports: bool = False,
    as_module: bool = False,
) -> None:
    """Automatically profile a script and run it.

    Profile functions, classes & modules specified in prof_mod without needing to add
    @profile decorators.

    Args:
        script_file (str):
            path to script being profiled.

        ns (dict):
            "locals" from kernprof scope.

        prof_mod (List[str]):
            list of imports to profile in script.
            passing the path to script will profile the whole script.
            the objects can be specified using its dotted path or full path (if applicable).

        profile_imports (bool):
            if True, when auto-profiling whole script, profile all imports aswell.

        as_module (bool):
            Whether we're running script_file as a module
    """

    class restore_dict:
        def __init__(self, d: MutableMapping[str, Any], target=None):
            self.d = d
            self.target = target
            self.copy: Mapping[str, Any] | None = None

        def __enter__(self):
            assert self.copy is None
            self.copy = dict(self.d)
            return self.target

        def __exit__(self, *_, **__):
            self.d.clear()
            if self.copy is not None:
                self.d.update(self.copy)
            self.copy = None

    Profiler: type[AstTreeModuleProfiler] | type[AstTreeProfiler]
    namespace: MutableMapping[str, Any]
    ctx: ContextManager

    if as_module:
        Profiler = AstTreeModuleProfiler
        module_name = modpath_to_modname(script_file)
        if not module_name:
            raise ModuleNotFoundError(
                f'script_file = {script_file!r}: cannot find corresponding module'
            )

        module_obj = types.ModuleType(module_name)
        namespace = vars(module_obj)
        namespace.update(ns)

        # Set the `__spec__` correctly
        module_obj.__spec__ = importlib.util.find_spec(module_name)

        # Set the module object to `sys.modules` via a callback, and
        # then restore it via the context manager
        callback = functools.partial(
            operator.setitem, sys.modules, '__main__', module_obj
        )
        ctx = restore_dict(sys.modules, callback)
    else:
        Profiler = AstTreeProfiler
        namespace = ns
        ctx = contextlib.nullcontext(lambda: None)

    profiler = Profiler(script_file, prof_mod, profile_imports)
    tree_profiled = profiler.profile()

    _extend_line_profiler_for_profiling_imports(ns[PROFILER_LOCALS_NAME])
    code_obj = compile(tree_profiled, script_file, 'exec')
    with ctx as callback:
        callback()
        exec(code_obj, cast(Dict[str, Any], namespace), namespace)
