"""Manages the rocm-sdk-devel package.

The devel package is special in some key ways:

* Since it contains distribution (wheel) unsafe files like symlinks, it is
  distributed under the `rocm_sdk_devel` package as a `_devel.tar` or
  `_devel.tar.xz` file that is intended to be expanded on use.
* This tarball is intended to be expanded into the site-lib directory that
  contains the ROCM distribution packages and will result in a top-level
  python package named like `_rocm_sdk_devel_linux_x86_64` that is a sibling
  to other packages like `_rocm_sdk_core_linux_x86_64`.
* For any files already contained in one of the runtime packages, a relative
  symlink to the correct sibling will be stored.
* Any files not in one of the runtime packages will be included verbatim in the
  tarball.
* RPATH setup relies on this sibling behavior and is already encoded properly
  in the runtime packages.

In order to make this work, we dynamically extend the distribution package on
use, modifying the dist-info RECORD file to include all newly expanded files in
accordance with the PyPA documentation:
  https://packaging.python.org/en/latest/specifications/recording-installed-packages/
Note that this puts us in the category of creating a self-modifying package,
which is strongly discouraged but not prohibited. We deem the tradeoff worth
it, as the alternative is to increase the package size by 2-5x and break
symlink relationships.
"""

import importlib.metadata as md
from pathlib import Path
import shutil
import tarfile

from . import _dist_info as di


def get_devel_root() -> Path:
    try:
        import rocm_sdk_devel
    except ModuleNotFoundError as e:
        raise ModuleNotFoundError(
            "ROCm SDK development package is not installed. This can typically be "
            "obtained by installing `rocm-sdk[devel]` from your package manager"
        ) from e
    rocm_sdk_devel_path = _get_package_path(rocm_sdk_devel)
    if rocm_sdk_devel_path is None:
        raise ModuleNotFoundError(
            "rocm_sdk_devel expected to be defined by an __init__.py file"
        )
    site_lib_path = rocm_sdk_devel_path.parent
    devel_py_pkg_name = di.ALL_PACKAGES["devel"].get_py_package_name()
    devel_py_pkg_path = site_lib_path / devel_py_pkg_name
    if (devel_py_pkg_path / "__init__.py").exists():
        return devel_py_pkg_path

    _expand_devel_contents(rocm_sdk_devel_path, site_lib_path)
    if not (devel_py_pkg_path / "__init__.py").exists():
        raise ImportError(
            f"Expanding {devel_py_pkg_name} did not produce a valid Python package"
        )
    return devel_py_pkg_path


# Gets the path of a module presumed to be a package defined by an __init__.py
# file. Returns None if it is a namespace package or another kind of module.
def _get_package_path(m) -> Path | None:
    if m.__file__ is None:
        return None
    p = Path(m.__file__)
    if p.name == "__init__.py":
        return p.parent  # Directory containing __init__.py
    return None


def _expand_devel_contents(rocm_sdk_devel_path: Path, site_lib_path: Path):
    # Resolve the Python package to its distribution package name and find the
    # RECORD file.
    dist_names = md.packages_distributions()["rocm_sdk_devel"]
    assert len(dist_names) == 1  # Would only be != 1 for namespace package.
    dist_name = dist_names[0]
    dist_files = md.files(dist_name)
    if dist_files is None:
        raise ImportError(
            "Cannot expand the `rocm-sdk[devel]` package because it was not installed "
            "by a user-mode package manager and is managed by the system. Please "
            "install the `rocm-sdk` in a virtual environment."
        )
    for record_pkg_file in dist_files:
        if record_pkg_file.name == "RECORD" and record_pkg_file.parent.name.endswith(
            ".dist-info"
        ):
            break
    else:
        raise ImportError(
            f"No distribution RECORD found for the `{dist_name}` distribution package."
        )

    # Resolve to a physical file.
    record_path = record_pkg_file.locate()

    # Find the tarfile.
    tarfile_path = rocm_sdk_devel_path / "_devel.tar.xz"
    if tarfile_path.exists():
        tarfile_mode = "r:xz"
    else:
        tarfile_path = rocm_sdk_devel_path / "_devel.tar"
        if not tarfile_path.exists():
            raise ImportError(
                f"Expected to find _devel.tar or _devel.tar.xz in {rocm_sdk_devel_path}"
            )
        tarfile_mode = "r"

    dist_file_path_names = [str(df) for df in dist_files]
    _lock_and_expand(
        site_lib_path,
        tarfile_path,
        tarfile_mode,
        record_path,
        dist_file_path_names,
    )


def _lock_and_expand(
    site_lib_path: Path,
    tarfile_path: Path,
    tarfile_mode: str,
    record_path: Path,
    dist_file_path_names: set[str],
):
    # When extracting, we note the directory paths of each entry and on the first
    # access, clean it up if it is already present. This works around package manager
    # races where in certain uninstall situations, some amount of the directory tree
    # may not be fully removed (this presently happens with dangling symlinks).
    # Cleaning it ensures consistent re-install behavior.
    clean_dir_paths: set[Path] = set()

    def _clean_dir(dir: Path):
        clean_dir_paths.add(dir)
        if dir.exists():
            shutil.rmtree(dir, ignore_errors=False)

    with open(record_path, "at") as record_file:
        _lock_file(record_file)
        try:
            with tarfile.open(tarfile_path, tarfile_mode) as tf:
                while ti := tf.next():
                    dest_path = site_lib_path / ti.name
                    if ti.isfile() or ti.issym():
                        parent_path = dest_path.parent
                        if parent_path not in clean_dir_paths:
                            _clean_dir(parent_path)
                        tf.extract(ti, path=site_lib_path)
                        if ti.name not in dist_file_path_names:
                            # CSV record:
                            #   path
                            #   hash (empty)
                            #   size (empty)
                            record_file.write(f"{ti.name},,\n")
                    elif ti.isdir():
                        # We don't generally have directory entries, but handle
                        # them if we do.
                        if dest_path not in clean_dir_paths:
                            _clean_dir(dest_path)
                        tf.extract(ti, path=site_lib_path)
            tarfile_path.unlink()
        finally:
            _unlock_file(record_file)


def _lock_file(f):
    # TODO: This is Posix specific file locking. We should fork on Windows to use
    # msvcrt locking.
    import fcntl, os

    fcntl.lockf(f, fcntl.LOCK_EX)


def _unlock_file(f):
    # TODO: This is Posix specific file locking. We should fork on Windows to use
    # msvcrt locking.
    import fcntl, os

    fcntl.lockf(f, fcntl.LOCK_UN)
