"""Create a reviewed environment-only Docker context from a fixed Conda prefix."""
import argparse
import hashlib
import importlib.metadata
import io
import json
import os
import shutil
import tarfile
from pathlib import Path


WEIGHT_SUFFIXES = {".pt", ".ckpt", ".safetensors", ".onnx", ".model"}
SKIP_DIRS = {"__pycache__", ".git", ".cache", "conda-meta", "tests", "test"}
SKIP_NAMES = {"direct_url.json", "pip.conf", ".netrc", ".condarc", "hfkey.txt", "token", "stored_tokens", "credentials", "id_rsa", "id_ed25519", ".npmrc", ".pypirc"}
REVIEWED_BINARY_RESOURCES = {
    "matplotlib/mpl-data/sample_data/goog.npz", "matplotlib/mpl-data/sample_data/jacksboro_fault_dem.npz",
    "matplotlib/mpl-data/sample_data/topobathy.npz", "scipy/stats/_sobol_direction_numbers.npz",
    "sentencepiece/package_data/nfkc.bin", "sentencepiece/package_data/nfkc_cf.bin",
    "sentencepiece/package_data/nmt_nfkc.bin", "sentencepiece/package_data/nmt_nfkc_cf.bin",
    "text_unidecode/data.bin", "whisper/assets/mel_filters.npz",
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--prefix", type=Path, default=Path("/opt/conda/envs/hybrid"))
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    prefix = args.prefix.resolve()
    assert prefix == Path("/opt/conda/envs/hybrid"), "Snapshot preserves its original prefix"
    args.out.mkdir(parents=True, exist_ok=False)
    files, excluded = [], []
    for directory, dirs, names in os.walk(prefix, followlinks=False):
        dirs[:] = sorted(d for d in dirs if d not in SKIP_DIRS)
        for d in dirs:
            path = Path(directory) / d
            if path.is_symlink():
                assert path.resolve().is_relative_to(prefix), f"External directory link: {path}"
                files.append(path)
        for name in sorted(names):
            path = Path(directory) / name
            rel = path.relative_to(prefix)
            if name in SKIP_NAMES or name.startswith(".env") or path.suffix in WEIGHT_SUFFIXES or path.suffix == ".pyc":
                excluded.append(str(rel))
                continue
            if path.suffix == ".pth" and path.parent.name != "site-packages":
                excluded.append(str(rel))
                continue
            if path.suffix.lower() in {".bin", ".pkl", ".npz", ".h5", ".hdf5"}:
                resource = str(rel).removeprefix("lib/python3.12/site-packages/")
                if resource not in REVIEWED_BINARY_RESOURCES:
                    raise ValueError(f"Review potential model/data before including: {rel}")
            if path.is_symlink():
                assert path.resolve().is_relative_to(prefix), f"External file link: {path}"
            files.append(path)
    manifest = []
    patch_prefix = b"import inspect\nif not hasattr(inspect, 'getargspec'):\n    inspect.getargspec = inspect.getfullargspec\n"
    with tarfile.open(args.out / "environment.tar", "x", dereference=False) as archive:
        for path in sorted(set(files)):
            rel = path.relative_to(prefix)
            info = archive.gettarinfo(str(path), arcname=str(rel))
            info.uid = info.gid = 0
            info.uname = info.gname = "root"
            if info.issym():
                archive.addfile(info)
                manifest.append({"file": str(rel), "link": info.linkname})
                continue
            if not info.isfile():
                raise ValueError(f"Unsupported file: {rel}")
            modified = str(rel) == "lib/python3.12/site-packages/chumpy/__init__.py"
            if modified:
                original = path.read_bytes()
                data = patch_prefix + original
                info.size = len(data)
                archive.addfile(info, io.BytesIO(data))
                digest = hashlib.sha256(data).hexdigest()
            else:
                h = hashlib.sha256()
                with path.open("rb") as source:
                    for block in iter(lambda: source.read(1024 * 1024), b""):
                        h.update(block)
                digest = h.hexdigest()
                with path.open("rb") as source:
                    archive.addfile(info, source)
            entry = {"file": str(rel), "bytes": info.size, "sha256": digest}
            if modified:
                entry["patch"] = "chumpy-local inspect.getargspec fallback for Python 3.12"
                entry["original_sha256"] = hashlib.sha256(original).hexdigest()
            manifest.append(entry)
    versions = sorted({f"{dist.metadata['Name']}=={dist.version}" for dist in importlib.metadata.distributions() if dist.metadata.get("Name")})
    (args.out / "requirements.snapshot.txt").write_text("\n".join(versions) + "\n")
    report = {
        "prefix": str(prefix), "snapshot_kind": "installed dependencies, not a pip-rebuild lock",
        "excluded": excluded, "skipped_directories": sorted(SKIP_DIRS), "files": manifest,
        "scope": "No project, user home, HF cache, datasets, pretrained weights or credentials; optional package test/LPIPS weights removed",
    }
    (args.out / "environment-inventory.json").write_text(json.dumps(report, indent=2))
    source = Path(__file__).resolve().parent
    for name in ("Dockerfile", "smoke_test.py"):
        shutil.copy2(source / name, args.out / name)
    print(json.dumps({"context": str(args.out), "tar_bytes": (args.out / "environment.tar").stat().st_size, "entries": len(manifest), "excluded_files": len(excluded)}), flush=True)


if __name__ == "__main__":
    main()