"""A collection of functions which are triggered automatically by finder when
tensorflow package is included.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from cx_Freeze._compat import IS_MINGW, IS_WINDOWS
from cx_Freeze.module import Module, ModuleHook

if TYPE_CHECKING:
    from cx_Freeze.finder import ModuleFinder


__all__ = ["Hook"]


class Hook(ModuleHook):
    """The Hook class for tensorflow."""

    def tensorflow(self, finder: ModuleFinder, module: Module) -> None:
        """Hook for tensorflow. Tested in Windows and Linux."""
        module_path = module.file.parent
        site_packages_path = module_path.parent

        # implicitly loaded packages
        finder.include_package("tensorboard")
        finder.include_package("tensorflow.compiler")
        finder.include_package("tensorflow.python")

        # support for plugins
        source_path = site_packages_path / "tensorflow-plugins"
        if source_path.is_dir():
            pattern = "*.dll" if (IS_WINDOWS or IS_MINGW) else "*.so"
            for source in source_path.rglob(pattern):  # type: Path
                target = "lib" / source.relative_to(site_packages_path)
                finder.include_files(source, target)

        # patch the code to search the correct directory
        code_string = module.file.read_text(encoding="utf_8")
        code_string = code_string.replace(
            "_site_packages_dirs = []",
            "_site_packages_dirs = [_os.path.join(_sys.prefix, 'lib')]",
        )
        code_string = code_string.replace(
            "_current_file_location = ",
            "_current_file_location = __file__.replace('library.zip', '.')  #",
        )
        module.code = compile(
            code_string,
            module.file.as_posix(),
            "exec",
            dont_inherit=True,
            optimize=finder.optimize,
        )

        # installed version of tensorflow is a variant?
        if module.distribution is None:
            for name in (
                "tensorflow-aarch64",
                "tensorflow-cpu",
                "tensorflow-cpu-aws",
                "tensorflow-gpu",
                "tensorflow-intel",
                "tensorflow-macos",
                "tensorflow-rocm",
            ):
                module.update_distribution(name)
                if module.distribution is not None:
                    break

        # remove run-time warning
        # WARNING:tensorflow:AutoGraph is not available in this environment:...
        source_path = site_packages_path / "tensorflow/python/autograph"
        for source in source_path.rglob("*.py"):  # type: Path
            target = "lib" / source.relative_to(site_packages_path)
            finder.include_files(source, target)
        finder.exclude_module("tensorflow.python.autograph")
