import subprocess
import json
from datetime import datetime
import re
import os
import argparse

node_dict = {
    # "c0f6c83c-01.cloud.together.ai": "gpu01",
    # "c0f6c83c-02.cloud.together.ai": "gpu02",
    # "c0f6c83c-03.cloud.together.ai": "gpu03",
    # "c0f6c83c-04.cloud.together.ai": "gpu04",
    # "c0f6c83c-05.cloud.together.ai": "gpu05",
    # "c0f6c83c-06.cloud.together.ai": "gpu06",
    # "c0f6c83c-07.cloud.together.ai": "gpu07",
    # "c0f6c83c-08.cloud.together.ai": "gpu08",
    # "c0f6c83c-09.cloud.together.ai": "gpu09",
    # "c0f6c83c-10.cloud.together.ai": "gpu10",
    # "c0f6c83c-11.cloud.together.ai": "gpu11",
    # "c0f6c83c-12.cloud.together.ai": "gpu12",
    # "c0f6c83c-13.cloud.together.ai": "gpu13",
    # "c0f6c83c-14.cloud.together.ai": "gpu14",
    # "c0f6c83c-15.cloud.together.ai": "gpu15",
    # "c0f6c83c-16.cloud.together.ai": "gpu16",
    # "c0f6c83c-17.cloud.together.ai": "gpu17",
    # "c0f6c83c-18.cloud.together.ai": "gpu18",
    # "c0f6c83c-19.cloud.together.ai": "gpu19",
    # "c0f6c83c-20.cloud.together.ai": "gpu20",
    "c0f6c83c-21.cloud.together.ai": "gpu21",
    "c0f6c83c-22.cloud.together.ai": "gpu22",
    "c0f6c83c-23.cloud.together.ai": "gpu23",
    "c0f6c83c-24.cloud.together.ai": "gpu24",
    "c0f6c83c-25.cloud.together.ai": "gpu25",
    "c0f6c83c-26.cloud.together.ai": "gpu26",
    "c0f6c83c-27.cloud.together.ai": "gpu27",
    "c0f6c83c-28.cloud.together.ai": "gpu28",
    "c0f6c83c-33.cloud.together.ai": "gpu33",
    "c0f6c83c-34.cloud.together.ai": "gpu34",
    "c0f6c83c-35.cloud.together.ai": "gpu35",
    "c0f6c83c-36.cloud.together.ai": "gpu36",
    "c0f6c83c-37.cloud.together.ai": "gpu37",
    "c0f6c83c-38.cloud.together.ai": "gpu38",
}

## 06 07 11 12 13 15 18 20 21 25 26 

def run_ssh_command(node_alias, command):
    try:
        result = subprocess.run(
            ["ssh", node_alias, command],
            capture_output=True,
            text=True,
            timeout=30
        )
        if result.returncode != 0:
            print(f"[{node_alias}] Command error: {result.stderr.strip()}")
            return ""
        return result.stdout.strip()
    except Exception as e:
        print(f"[{node_alias}] SSH command exception: {e}")
        return ""


def get_gpu_info(node_alias):
    output = run_ssh_command(node_alias, "nvidia-smi")
    gpu_info = []

    mem_matches = re.findall(r'(\d+)MiB\s*/\s*(\d+)MiB', output)
    power_matches = re.findall(r'(\d+)W\s*/\s*(\d+)W', output)

    for i, ((used_mem, total_mem), (used_power, max_power)) in enumerate(zip(mem_matches, power_matches)):
        gpu_info.append({
            'gpu_id': i,
            'used_memory_mb': int(used_mem),
            'total_memory_mb': int(total_mem),
            'used_memory_gb': round(int(used_mem) / 1024, 2),
            'total_memory_gb': round(int(total_mem) / 1024, 2),
            'power_usage_w': int(used_power),
            'power_limit_w': int(max_power),
            'power_usage_percent': round((int(used_power) / int(max_power)) * 100, 2)
        })

    return gpu_info


def get_system_resources(node_alias):
    top_output = run_ssh_command(node_alias, "top -bn1")
    free_output = run_ssh_command(node_alias, "free -b")

    cpu_idle = re.search(r'%Cpu\(s\):.+?(\d+\.\d+)\s+id', top_output)
    cpu_idle_percent = float(cpu_idle.group(1)) if cpu_idle else 0.0

    memory_info = free_output.split('\n')
    mem_line = [line for line in memory_info if 'Mem:' in line][0].split()

    total_memory = int(mem_line[1]) / (1024 ** 3)
    used_memory = int(mem_line[2]) / (1024 ** 3)
    free_memory = int(mem_line[3]) / (1024 ** 3)
    available_memory = int(mem_line[6]) / (1024 ** 3)

    return {
        'cpu_idle_percent': cpu_idle_percent,
        'memory_total_gb': round(total_memory, 2),
        'memory_used_gb': round(used_memory, 2),
        'memory_free_gb': round(free_memory, 2),
        'memory_available_gb': round(available_memory, 2),
        'memory_free_percent': round((free_memory / total_memory) * 100, 2),
        'memory_available_percent': round((available_memory / total_memory) * 100, 2)
    }


def main(node_alias, node_id, save_folder_dir):
    try:
        gpu_info = get_gpu_info(node_alias)
        sys_resources = get_system_resources(node_alias)

        monitoring_data = {
            'timestamp': datetime.now().isoformat(),
            'gpu_info': gpu_info,
            'system_resources': sys_resources
        }

        print(f"\n####### Node: {node_alias} ({node_id}) #######")
        for gpu in gpu_info:
            print(f"GPU {gpu['gpu_id']}: {gpu['used_memory_gb']} GB used / {gpu['total_memory_gb']} GB")
        print(f"CPU idle %: {sys_resources['cpu_idle_percent']}")
        print(f"Memory free %: {sys_resources['memory_free_percent']}")

        os.makedirs(save_folder_dir, exist_ok=True)
        filename = f"{save_folder_dir}/{node_id}.json"
        with open(filename, 'w') as f:
            json.dump(monitoring_data, f, indent=4)
        print(f"Saved to {filename}")

    except Exception as e:
        print(f"[{node_alias}] Monitoring failed: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--save_dir', type=str, default='tmp_data',
                        help='Base directory to save monitoring data')
    args = parser.parse_args()

    timestamp_dir = f"{args.save_dir}/together_monitoring_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    for node_alias, node_id in node_dict.items():
        main(node_alias, node_id, timestamp_dir)