# Sets up interactive sessions in ipython to automatically import missing modules using nix Depends
# on nix-index/nix-locate being installed and functional, see
# <https://github.com/nix-community/nix-index-database> for a prebuilt database.
#
# How it works:
# - During startup, the sysconfig is used to determine where to find site-packages and the probable
#   nixpkgs variable name for the python scope.
# - To find a package, nix-locate is used to try to find a package which exports
#   {sitePackages}/{name}/__init__.py. If that is found, it attempts to download nixpkgs#{result}.
# - If that fails, it instead attempts to build nixpkgs#{pythonScope}.{name}.
# - When either of those strategies result in nix store paths, the store paths are recursively
#   processed (taking propagated-build-inputs into account) and any {sitePackages} found in any
#   dependency is added to sys.path.
# - Finally, the import is tried again with importlib.machinery.PathFinder
# - A directory is created under $XDG_RUNTIME_DIR/ipython per process which contains GC roots for
#   any packages loaded this way. When ipython exits, the per-process directory is deleted. This
#   ensures that a GC during a running ipython session doesn't nuke any temporarily installed
#   packages.
#
# Limitations:
# - The only way to be able to look up a package that is not named the same as its top-level module
#   is using nix-locate. The index is only built for the default python scope, and not any other
#   versions of python. This limitation could be worked around by scanning for pythonImportsCheck in
#   nixpkgs but this is not performed currently. Therefore, eg `import Crypto` works on the default
#   python version for a given NixOS installation, but not any other version. This script also
#   provids install_with_nix(package), which can be called with eg install_with_nix("pycryptodome"),
#   after which such imports should work.
# - Support for propagated-build-inputs is very rudimentary. For example, a python package which
#   depends on a binary being present in PATH which is declared in propagatedBuildInputs would not
#   currently work. However, pure-python dependencies work as expected.
# - This will probably explode and should not be used outside of **interactive** sessions, there is
#   a check to make sure an ipython input is in the call stack when the MetaPathFinder is triggered.


def _nix_importer_init():
    import atexit
    import importlib.abc
    import importlib.machinery
    import importlib.util
    import inspect
    import os
    import pathlib
    import shutil
    import subprocess
    import sys
    import sysconfig
    import tempfile
    from types import ModuleType
    from typing import Sequence


    runtime_dir = pathlib.Path(
            os.environ.get("XDG_RUNTIME_DIR",
                           str(pathlib.Path("/run/user") / str(os.getuid()))))
    ipython_dir = runtime_dir / "ipython"
    if not ipython_dir.is_dir():
        ipython_dir.mkdir()
    session_dir = ipython_dir / f"session-{os.getpid()}"
    if session_dir.exists():
        shutil.rmtree(session_dir)
    session_dir.mkdir()
    atexit.register(lambda: shutil.rmtree(session_dir))


    class NixImporter(importlib.abc.MetaPathFinder):
        def __init__(self) -> None:
            self._result_counter = 0
            self._subfinder = importlib.machinery.PathFinder()

            config_vars = sysconfig.get_config_vars()
            config_vars["base"] = ""
            path = sysconfig.get_path("purelib", scheme="posix_prefix", vars=config_vars)
            # remove leading / if it exists
            self._base_path = ("/" / pathlib.Path(path)).relative_to("/")
            self._base_packages = f"python{config_vars['py_version_nodot']}Packages"


        def find_spec(self, fullname: str,
                      path: Sequence[str] | None,
                      target: None | ModuleType = None) -> importlib.machinery.ModuleSpec | None:
            # Deny autoimport to the nix store, and only allow it if the traceback indicates ipython
            f = inspect.currentframe()
            while f is not None:
                if f.f_code.co_filename.startswith("/nix/store"):
                    return None
                elif f.f_code.co_filename.startswith("<ipython-input"):
                    break
                f = f.f_back
            else:
                return None

            find_path = self._base_path
            for element in fullname.split("."):
                find_path = find_path / element
            find_path = find_path / "__init__.py"

            print(f"Module {fullname} not found, searching NixOS...")
            command = ["nix-locate", "-w", "--minimal", str(find_path)]

            selected_target: str | None = None

            try:
                output = subprocess.check_output(command).strip().decode()
                if len(output) > 0:
                    if len(matches := output.split("\n")) > 1:
                        print("Multiple matches found!", matches)
                        return None
                    else:
                        selected_target = output
            except Exception:
                print("Failed to search")

            if selected_target is None and "." not in fullname:
                selected_target = f"{self._base_packages}.{fullname}"

            if selected_target is None:
                print("Could not find installation target!")
                return None

            try:
                print("Installing", selected_target)
                self._install_with_nix(selected_target)
                return self._subfinder.find_spec(fullname, path, target)
            except Exception:
                print("Failed to install")

            return None


        def invalidate_caches(self) -> None:
            self._subfinder.invalidate_caches()

        def _process_propagated_build_inputs(self, path: pathlib.Path) -> None:
            python_path = path / self._base_path
            if python_path.exists() and str(python_path) not in sys.path:
                print("Appending path", str(python_path))
                sys.path.append(str(python_path))

            nix_support = path / "nix-support"
            if not nix_support.is_dir():
                return
            propagated_build_inputs = nix_support / "propagated-build-inputs"
            if not propagated_build_inputs.is_file():
                return
            with propagated_build_inputs.open("r") as f:
                paths = f.read().strip().split(" ")
                for child_path in paths:
                    self._process_propagated_build_inputs(pathlib.Path(child_path))

        def _install_with_nix(self, path: str) -> None:
            if not path.startswith("/nix/store"):
                path = "nixpkgs#" + path

            out_paths = subprocess.check_output(
                    ["nix", "build", "--out-link", f"link-{self._result_counter}", "--print-out-paths", path],
                    cwd=session_dir).strip().decode().split("\n")
            self._result_counter += 1

            for result in out_paths:
                self._process_propagated_build_inputs(pathlib.Path(result))


    importer = NixImporter()
    sys.meta_path.append(importer)

    return importer


install_with_nix = _nix_importer_init()._install_with_nix
del _nix_importer_init