#!/usr/bin/env bash
# SUM Parts - repair packages truncated by the WSL VM crash
#
# The WSL instance died (Wsl/Service/CreateInstance/E_FAIL) partway through the
# dependency install. Everything pip had open at that moment landed on disk as
# 0-byte files. The failure mode is nasty because import still succeeds -- the
# module is simply empty, so you get things like
#
#     from torch_scatter import scatter
#     TypeError: 'module' object is not callable
#
# pointing at the call site rather than at the broken install.
#
# This reinstalls the batch that was in flight. Zero-byte __init__.py files are
# normal (namespace markers), so the scan below ignores those.
set -euo pipefail

CONDA_ROOT="$HOME/miniconda3"
ENV_NAME="sumparts"

source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"

SITE=$(python -c "import site; print(site.getsitepackages()[0])")

scan() {
    find "$SITE" -type f \( -name '*.py' -o -name '*.so' \) -size 0 \
        ! -name '__init__.py' | wc -l
}

echo "=== suspicious 0-byte files before: $(scan) ==="

# The dependency batch that was installing when the VM went down.
PKGS=(
    plyfile scikit-learn easydict PyYAML tensorboard termcolor tqdm
    multimethod h5py matplotlib pandas shortuuid gdown Cython pyvista
    wandb trimesh pillow
)

echo "=== force-reinstalling ${#PKGS[@]} packages ==="
pip install --no-cache-dir --force-reinstall "${PKGS[@]}"

# Pins that other steps depend on; --force-reinstall above can pull them up.
echo "=== restoring pins ==="
pip install --no-cache-dir "numpy<2" "setuptools==69.5.1" "ninja==1.11.1.1"

# torch-scatter must match the exact torch build, so it needs its own index.
python -c "import torch_scatter, inspect; assert callable(torch_scatter.scatter)" 2>/dev/null || {
    echo "=== reinstalling torch-scatter ==="
    pip install --no-cache-dir --force-reinstall torch-scatter \
        -f https://data.pyg.org/whl/torch-2.0.1+cu118.html
}

echo
echo "=== suspicious 0-byte files after: $(scan) ==="
find "$SITE" -type f \( -name '*.py' -o -name '*.so' \) -size 0 \
    ! -name '__init__.py' | head -20

echo
echo "=== import check ==="
python - <<'PY'
import importlib
mods = ["numpy", "torch", "torch_scatter", "plyfile", "sklearn", "matplotlib",
        "pandas", "h5py", "yaml", "easydict", "tqdm", "termcolor", "multimethod",
        "wandb", "trimesh", "PIL", "gdown"]
bad = []
for m in mods:
    try:
        importlib.import_module(m)
        print(f"{m:14s} OK")
    except Exception as e:
        bad.append(m)
        print(f"{m:14s} FAIL {type(e).__name__}: {e}")
import torch_scatter
print("torch_scatter.scatter callable:", callable(torch_scatter.scatter))
raise SystemExit(1 if bad or not callable(torch_scatter.scatter) else 0)
PY

echo "REPAIR DONE"