"""Check the CUDA/APE runtime needed by Orienter Stage 2."""

import argparse
import importlib
import importlib.metadata
import json
import platform
import sys


HISTORICAL_REFERENCE = {
    "python": "3.9.18",
    "torch": "2.2.0+cu121",
    "torchvision": "0.17.0+cu121",
    "detectron2": "0.6",
    "detrex": "0.3.0",
}


def _version(module_name, distribution_name=None):
    module = importlib.import_module(module_name)
    value = getattr(module, "__version__", None)
    if value is not None:
        return str(value)
    return importlib.metadata.version(distribution_name or module_name)


def check_environment(require_cuda=True, require_ape_extension=True):
    report = {
        "python": platform.python_version(),
        "historical_reference": dict(HISTORICAL_REFERENCE),
        "versions": {},
        "cuda_available": False,
        "ape_extension": False,
        "errors": [],
    }

    for module_name, distribution_name in (
        ("torch", "torch"),
        ("torchvision", "torchvision"),
        ("detectron2", "detectron2"),
        ("detrex", "detrex"),
    ):
        try:
            report["versions"][module_name] = _version(module_name, distribution_name)
        except Exception as exc:
            report["errors"].append(f"cannot import {module_name}: {type(exc).__name__}: {exc}")

    try:
        torch = importlib.import_module("torch")
        report["torch_cuda_runtime"] = torch.version.cuda
        report["cuda_available"] = bool(torch.cuda.is_available())
        if require_cuda and not report["cuda_available"]:
            report["errors"].append("torch.cuda.is_available() is false")
    except Exception:
        pass

    if require_ape_extension:
        try:
            importlib.import_module("ape._C")
            report["ape_extension"] = True
        except Exception as exc:
            report["errors"].append(
                f"cannot import compiled ape._C extension: {type(exc).__name__}: {exc}"
            )

    return report


def build_parser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--allow-no-cuda", action="store_true")
    parser.add_argument("--skip-ape-extension", action="store_true")
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    report = check_environment(
        require_cuda=not args.allow_no_cuda,
        require_ape_extension=not args.skip_ape_extension,
    )
    print(json.dumps(report, indent=2, sort_keys=True))
    return 1 if report["errors"] else 0


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