import subprocess
import pandas as pd
import tqdm


def get_node_addrs():
    with open("/etc/hosts", "r") as fi:
        res = fi.readlines()
    res = [r.strip() for r in res]
    res = [r for r in res if r]
    ii = res.index("# END ANSIBLE MANAGED BLOCK BASTION")
    assert res[ii + 1].startswith("# BEGIN ANSIBLE MANAGED BLOCK")
    hosts = res[ii + 2 :]
    hosts = [h.split()[1].replace("-rdma.local.rdma", "") for h in hosts if not h.startswith("#")]
    hosts = list(set(hosts))
    assert len(hosts) <= 30
    return hosts


def parse_line(line):
    items = line.split(",")
    items = [i.replace("%", "").strip() for i in items]
    return [float(i) for i in items]


def parse_header(header):
    header = header.split(", ")
    header = [h.replace("[%]", "").strip() for h in header]
    return header


def run_smi_cmd_on_node(addr):
    fmt_str = "utilization.memory,utilization.gpu"
    fmt_str = f"--query-gpu={fmt_str}"
    cmd = ["ssh", addr, "nvidia-smi", fmt_str, "--format=csv"]

    ret = subprocess.check_output(cmd).decode("utf-8")
    lines = ret.strip().split("\n")
    header, lines = lines[0], lines[1:]
    header = parse_header(header)
    lines = [parse_line(line) for line in lines]
    df = pd.DataFrame(lines, columns=header)
    df["node_addr"] = addr
    return df


if __name__ == "__main__":
    pieces = []
    node_addrs = get_node_addrs()
    for addr in tqdm.tqdm(node_addrs):
        pieces.append(run_smi_cmd_on_node(addr))
    df = pd.concat(pieces)
    gb = df.groupby("node_addr").mean()
    gb.index = [i.replace("compute-permanent-", "") for i in gb.index]
    gb.columns = [col + " (%)" for col in gb.columns]
    print(gb)