"""Copy installed distributions between Python environments, file for file.

Why this exists rather than `pip install`: pip on the H20 host stalls indefinitely on PyPI
(the wheel for `posix_ipc` never finishes downloading), and the packages we need are already
installed and pinned in a reference image. Copying is exact and takes seconds.

Two things it handles that a naive `cp -r` does not:

* A distribution's files are not always under its own name. `pillow` installs `PIL` and
  `pillow.libs`, `pyzmq` installs `zmq` and `pyzmq.libs`. The file list comes from the
  distribution's metadata, so those sibling directories come along.
* Some packages must *not* be copied even when present, because the target image ships a
  version linked against libs the source does not have. `pillow` is the case in hand: copying
  11.3 into an image whose `pillow.libs` holds a different libtiff set leaves `PIL.Image`
  unable to import. Pass it in `--keep-system` and the target's own copy is reused.

    python scripts/copy_site_packages.py <src-site-packages> <dst-site-packages> \\
        --packages einops numpy pillow \\
        --keep-system pillow
"""

import argparse
import importlib.metadata as md
import shutil
import sys
from pathlib import Path

# Everything `requirements.txt` asks for that the engine imports at load time. torch,
# tilelang, tile-kernels and tvm-ffi are deliberately absent: they are either provided by
# the image or installed by setup_flagtree_image.sh from a reference container, because
# their builds are ABI-sensitive in ways a file copy only hides.
DEFAULT_PACKAGES = [
    "einops",
    "huggingface_hub",
    "ninja",
    "numpy",
    "orjson",
    "pillow",
    "posix_ipc",
    "psutil",
    "pybind11",
    "pyzmq",
    "requests",
    "safetensors",
    "scipy",
    "sympy",
    "tabulate",
    "tokenizers",
    "tqdm",
    "transformers",
    "xgrammar",
]


def owned_paths(dist: md.Distribution) -> list[Path]:
    """The top-level entries a distribution installs, plus its metadata directory."""
    root = Path(dist.locate_file(""))
    found: dict[str, Path] = {}

    # top_level.txt names the import roots, which is what we want when they differ from the
    # distribution name.
    try:
        declared = dist.read_text("top_level.txt") or ""
    except (FileNotFoundError, UnicodeDecodeError):
        declared = ""
    for line in declared.splitlines():
        name = line.strip()
        if name and (root / name).exists():
            found[name] = root / name

    # Fall back to the file list; single-module distributions have no top_level entry.
    if not found:
        for f in dist.files or []:
            first = f.parts[0]
            if first.endswith((".dist-info", ".egg-info", ".data")) or first.startswith(".."):
                continue
            if (root / first).exists():
                found[first] = root / first

    found[dist._path.name] = dist._path.resolve()
    return list(found.values())


def install(src: Path, dst: Path, name: str, dist: md.Distribution) -> list[str]:
    copied = []
    for path in owned_paths(dist):
        target = dst / path.name
        if target.is_symlink() or target.is_file():
            target.unlink()
        elif target.is_dir():
            shutil.rmtree(target)
        if path.is_dir() and not path.is_symlink():
            shutil.copytree(path, target, symlinks=True)
        else:
            # `.so` files and metadata are files; copy_file keeps the mode bits, which the
            # extension modules need.
            shutil.copy2(path, target, follow_symlinks=False)
        copied.append(path.name)
    return copied


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("src", type=Path, help="source site-packages directory")
    parser.add_argument("dst", type=Path, help="destination site-packages directory")
    parser.add_argument("--packages", nargs="+", default=DEFAULT_PACKAGES)
    parser.add_argument(
        "--keep-system",
        nargs="*",
        default=["pillow"],
        help="packages to leave to the destination image, because their vendored libs differ",
    )
    args = parser.parse_args()

    for p in (args.src, args.dst):
        if not p.is_dir():
            print(f"not a directory: {p}", file=sys.stderr)
            return 2

    # Reading metadata from `src` without installing it: the finder consults sys.path, so the
    # source tree goes first for the life of this process.
    sys.path.insert(0, str(args.src))

    missing = []
    for name in args.packages:
        if name in args.keep_system:
            print(f"skip      {name} (keeping the destination image's copy)")
            continue
        try:
            dist = md.distribution(name)
        except md.PackageNotFoundError:
            missing.append(name)
            print(f"missing   {name}")
            continue
        copied = install(args.src, args.dst, name, dist)
        print(f"{dist.version:<10}{name}: {', '.join(copied)}")

    if missing:
        print(f"\nnot found in {args.src}: {', '.join(missing)}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())