#!/usr/bin/env python3
"""Hardware profiles shared by portable crossover launchers and audits."""
import os
import subprocess


PROFILES = {
    "timan107_gtx1080": {
        "allowed_physical_gpu_indices": [5, 7],
        "allowed_device_names": ["NVIDIA GeForce GTX 1080"],
    },
    "a6000": {
        "allowed_physical_gpu_indices": None,
        "allowed_device_names": ["NVIDIA RTX A6000"],
    },
}
DEFAULT_PROFILE = "timan107_gtx1080"


def profile_choices():
    return tuple(PROFILES)


def policy(profile, physical_gpu_index=None):
    if profile not in PROFILES:
        raise ValueError(f"unknown crossover hardware profile: {profile}")
    configured = PROFILES[profile]
    allowed = configured["allowed_physical_gpu_indices"]
    if allowed is None:
        if physical_gpu_index is None:
            raise ValueError(
                f"{profile} policy requires a resolved physical GPU index"
            )
        allowed = [int(physical_gpu_index)]
    return {
        "profile": profile,
        "allowed_physical_gpu_indices": list(allowed),
        "allowed_device_names": list(configured["allowed_device_names"]),
        "single_gpu_per_process": True,
    }


def _query_physical_gpus():
    query = subprocess.run(
        [
            "nvidia-smi",
            "--query-gpu=index,uuid,name,memory.total",
            "--format=csv,noheader,nounits",
        ],
        check=True,
        capture_output=True,
        text=True,
    ).stdout.splitlines()
    rows = {}
    for line in query:
        index, uuid, name, memory = [
            value.strip() for value in line.split(",", 3)
        ]
        rows[index] = {
            "uuid": uuid,
            "name": name,
            "total_memory_mib": int(memory),
        }
    return rows


def physical_gpu_report(profile, dry_run=False):
    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
    if dry_run:
        return {
            "hardware_profile": profile,
            "cuda_visible_devices": visible,
            "physical_gpu_index": None,
            "physical_gpu_uuid": None,
            "physical_gpu_name": None,
            "total_memory_mib": None,
        }
    if visible is None or not visible.isdigit():
        raise RuntimeError(
            "formal crossover jobs require CUDA_VISIBLE_DEVICES to contain "
            "one physical numeric GPU index"
        )
    physical_index = int(visible)
    hardware_policy = policy(profile, physical_index)
    if physical_index not in hardware_policy["allowed_physical_gpu_indices"]:
        raise RuntimeError(
            f"physical GPU {physical_index} is not allowed by {profile}"
        )
    rows = _query_physical_gpus()
    if visible not in rows:
        raise RuntimeError(f"physical GPU {visible} not found")
    row = rows[visible]
    if row["name"] not in hardware_policy["allowed_device_names"]:
        raise RuntimeError(
            f"{profile} requires one of "
            f"{hardware_policy['allowed_device_names']}, found {row['name']}"
        )
    return {
        "hardware_profile": profile,
        "cuda_visible_devices": visible,
        "physical_gpu_index": physical_index,
        "physical_gpu_uuid": row["uuid"],
        "physical_gpu_name": row["name"],
        "total_memory_mib": row["total_memory_mib"],
    }


def policy_for_report(profile, report):
    return policy(profile, report["physical_gpu_index"])


def assert_hardware_report(report, hardware_policy):
    assert report["hardware_profile"] == hardware_policy["profile"]
    assert (
        report["physical_gpu_index"]
        in hardware_policy["allowed_physical_gpu_indices"]
    )
    assert report["physical_gpu_uuid"]
    assert report["physical_gpu_name"] in hardware_policy[
        "allowed_device_names"
    ]
    assert report["cuda_visible_devices"] == str(
        report["physical_gpu_index"]
    )