"""
ROCm / CUDA health check.

Reports:
  - Backend (CUDA vs ROCm)
  - GPU count, names, VRAM
  - PyTorch version
  - Diffusers / Transformers / PEFT versions
  - Whether HSA_OVERRIDE_GFX_VERSION is needed (ROCm only)

Resilient: a single library import failure is reported but doesn't crash the check.
"""
from __future__ import annotations


def main():
    import torch
    print("=" * 60)
    print("Indic Heritage Studio v2 — Environment Check")
    print("=" * 60)

    print(f"\nPyTorch version: {torch.__version__}")
    print(f"CUDA available: {torch.cuda.is_available()}")

    if torch.cuda.is_available():
        backend = "ROCm" if (hasattr(torch.version, "hip")
                              and torch.version.hip) else "CUDA"
        print(f"Backend: {backend}")
        if backend == "ROCm":
            print(f"ROCm/HIP version: {torch.version.hip}")

        print(f"\nGPU count: {torch.cuda.device_count()}")
        total_vram = 0.0
        for i in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(i)
            vram_gb = props.total_memory / 1e9
            total_vram += vram_gb
            print(f"  GPU {i}: {props.name} ({vram_gb:.1f} GB)")
        print(f"Total VRAM across all GPUs: {total_vram:.1f} GB")

    # Library versions (resilient — one failure shouldn't crash the check)
    print("\nLibrary versions:")
    for lib in ("diffusers", "transformers", "peft", "accelerate",
                "gradio", "safetensors", "controlnet_aux", "cv2",
                "compel", "bitsandbytes", "datasets"):
        try:
            mod = __import__(lib)
            ver = getattr(mod, "__version__", "unknown")
            print(f"  {lib}: {ver}")
        except ImportError as e:
            print(f"  {lib}: NOT INSTALLED ({e})")
        except Exception as e:
            print(f"  {lib}: IMPORT ERROR ({type(e).__name__}: {e})")

    # Multi-GPU strategy
    if torch.cuda.is_available() and torch.cuda.device_count() >= 4:
        print("\n✓ Multi-GPU mode: pipelines will be pinned to dedicated GPUs.")
        print("  T2I → GPU 0 | Style → GPU 1 | I2V → GPU 2 | ControlNet → GPU 3")
        print("  Batch workers → GPU 4-7")
    elif torch.cuda.is_available() and torch.cuda.device_count() >= 2:
        print("\n⚠ Multi-GPU mode: limited GPUs — pipelines will share.")
    else:
        print("\n⚠ Single-GPU mode — pipelines will load/unload on demand.")

    # CUDA_VISIBLE_DEVICES warning
    import os
    cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
    if cvd:
        print(f"\n⚠ CUDA_VISIBLE_DEVICES is set to '{cvd}' — this restricts GPU visibility.")
        print(f"  Run: unset CUDA_VISIBLE_DEVICES  to see all GPUs.")

    print("\n" + "=" * 60)


if __name__ == "__main__":
    main()