#!/usr/bin/env python3
"""Check that every compiled extension and import the sumv2 training path needs
is actually present.

Extension import names do not match their directory names, which is an easy way
to waste an hour chasing a build that already succeeded:

    openpoints/cpp/pointnet2_batch -> pointnet2_batch_cuda
    openpoints/cpp/pointops        -> pointops_cuda
    openpoints/cpp/chamfer_dist    -> chamfer
    openpoints/cpp/emd             -> emd_cuda      (package name: emd_ext)
    openpoints/cpp/subsampling     -> openpoints.cpp.subsampling.grid_subsampling

Run from PointNeXt_bundle/ (or anywhere, if openpoints is importable).
"""
import os
import sys
from pathlib import Path


def add_openpoints_to_path() -> Path | None:
    """Put PointNeXt_bundle on sys.path.

    main.py does this itself with a hardcoded '../../', but this script lives
    outside the repo, so walk up from cwd (then from the default clone path)
    until a directory containing openpoints/ turns up.
    """
    candidates = [Path.cwd(), *Path.cwd().parents,
                  Path.home() / "sum-parts/semantic_segmentation/PointNeXt_bundle"]
    for c in candidates:
        if (c / "openpoints" / "__init__.py").exists():
            sys.path.insert(0, str(c))
            return c
    return None


MODULES = [
    "torch",
    "numpy",
    "pointnet2_batch_cuda",
    "pointops_cuda",
    "chamfer",
    "emd_cuda",
    "torch_scatter",
    "plyfile",
    "wandb",
    "trimesh",
]

FROM_IMPORTS = [
    ("openpoints.cpp.subsampling", "grid_subsampling"),
    ("openpoints.models", "build_model_from_cfg"),
    ("openpoints.dataset", "build_dataloader_from_cfg"),
]


def main() -> int:
    ok = True

    root = add_openpoints_to_path()
    print(f"openpoints root: {root or 'NOT FOUND'}")
    if root is None:
        ok = False

    try:
        import torch
        print(f"torch {torch.__version__} | cuda {torch.version.cuda} | "
              f"available {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"device: {torch.cuda.get_device_name(0)}")
    except Exception as e:  # noqa: BLE001
        print(f"torch import failed: {e}")
        return 1

    print()
    for name in MODULES:
        try:
            __import__(name)
            print(f"{name:32s} OK")
        except Exception as e:  # noqa: BLE001
            ok = False
            print(f"{name:32s} FAIL  {type(e).__name__}: {e}")

    print()
    for mod, attr in FROM_IMPORTS:
        try:
            m = __import__(mod, fromlist=[attr])
            getattr(m, attr)
            print(f"{mod + '.' + attr:32s} OK")
        except Exception as e:  # noqa: BLE001
            ok = False
            print(f"{mod + '.' + attr:32s} FAIL  {type(e).__name__}: {e}")

    print()
    print("ALL OK" if ok else "SOME CHECKS FAILED")
    return 0 if ok else 1


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