from __future__ import annotations

from dataclasses import dataclass

import ray


@dataclass(frozen=True)
class ClusterAvailable:
    total_available_gpus: float
    total_available_npus: float


def get_cluster_available() -> ClusterAvailable:
    # Align with verl's fail-fast check which uses ray._private.state.available_resources_per_node().
    # This is a best-effort internal API and may change with Ray versions.
    try:
        import ray._private.state  # type: ignore

        per_node = ray._private.state.available_resources_per_node()
    except Exception:
        # If we cannot fetch per-node resources, conservatively return 0.
        return ClusterAvailable(total_available_gpus=0.0, total_available_npus=0.0)

    total_gpu = 0.0
    total_npu = 0.0
    for _, info in per_node.items():
        if not isinstance(info, dict):
            continue
        total_gpu += float(info.get("GPU", 0) or 0)
        total_npu += float(info.get("NPU", 0) or 0)
    return ClusterAvailable(total_available_gpus=total_gpu, total_available_npus=total_npu)


def ensure_ray_connected() -> None:
    ray.init(address="auto", ignore_reinit_error=True, log_to_driver=False)

