#!/usr/bin/env python3
"""
Nvitop 데이터 수집 에이전트
각 서버에서 실행되어 nvitop 데이터를 수집하고 중앙 서버로 전송합니다.
"""

import argparse
import json
import time
import sys
import requests
from datetime import datetime
from typing import Dict, List, Any
import subprocess
import os

try:
    import pynvml
    PYNVML_AVAILABLE = True
except ImportError:
    PYNVML_AVAILABLE = False
    print("Warning: pynvml not available, using nvitop command instead")

try:
    import psutil
    PSUTIL_AVAILABLE = True
except ImportError:
    PSUTIL_AVAILABLE = False
    print("Warning: psutil not available, process info will be limited")


def get_nvitop_data() -> Dict[str, Any]:
    """nvitop Python API를 사용하여 GPU 정보를 수집합니다."""
    try:
        # nvitop Python API 사용 시도
        try:
            from nvitop import Device
            devices = Device.all()
            
            gpus = []
            for device in devices:
                # GPU 정보 수집
                try:
                    gpu_name = device.name()
                    if not gpu_name or gpu_name.strip() == '':
                        gpu_name = f'GPU {device.index}'
                except:
                    gpu_name = f'GPU {device.index}'
                
                gpu_info = {
                    'index': device.index,
                    'name': gpu_name,
                    'utilization': device.gpu_utilization(),
                    'memory_utilization': device.memory_utilization(),
                    'memory_used': device.memory_used() / 1024**2,  # MB
                    'memory_total': device.memory_total() / 1024**2,  # MB
                    'temperature': device.temperature(),
                    'power': device.power_usage() / 1000.0 if device.power_usage() else 0,  # W
                    'processes': []
                }
                
                # 프로세스 정보
                if PSUTIL_AVAILABLE:
                    try:
                        processes = device.processes()
                        for proc in processes:
                            try:
                                gpu_info['processes'].append({
                                    'pid': proc.pid,
                                    'name': proc.name(),
                                    'memory': proc.gpu_memory() / 1024**2,  # MB
                                    'username': proc.username()
                                })
                            except:
                                pass
                    except:
                        pass
                
                gpus.append(gpu_info)
            
            return {'gpus': gpus}
        except ImportError:
            # nvitop Python API가 없는 경우 pynvml 사용
            if PYNVML_AVAILABLE:
                return get_pynvml_data()
            else:
                return get_nvidia_smi_data()
    except Exception as e:
        print(f"Error getting nvitop data: {e}")
        # Fallback to pynvml or nvidia-smi
        if PYNVML_AVAILABLE:
            return get_pynvml_data()
        else:
            return get_nvidia_smi_data()


def parse_nvitop_text(text: str) -> Dict[str, Any]:
    """nvitop 텍스트 출력을 파싱합니다."""
    gpus = []
    lines = text.strip().split('\n')
    
    for line in lines:
        if 'GPU' in line or 'NVIDIA' in line:
            # 간단한 파싱 (실제로는 더 복잡한 파싱 필요)
            parts = line.split()
            if len(parts) > 0:
                gpus.append({
                    'index': len(gpus),
                    'name': ' '.join(parts[:3]) if len(parts) >= 3 else 'Unknown',
                    'utilization': 0,
                    'memory_used': 0,
                    'memory_total': 0,
                    'temperature': 0,
                    'power': 0
                })
    
    return {'gpus': gpus}


def get_pynvml_data() -> Dict[str, Any]:
    """pynvml을 사용하여 GPU 정보를 수집합니다."""
    if not PYNVML_AVAILABLE:
        return {}
    
    try:
        pynvml.nvmlInit()
        device_count = pynvml.nvmlDeviceGetCount()
        gpus = []
        
        for i in range(device_count):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            
            # GPU 이름
            name = pynvml.nvmlDeviceGetName(handle).decode('utf-8')
            
            # 사용률
            util = pynvml.nvmlDeviceGetUtilizationRates(handle)
            
            # 메모리
            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            
            # 온도
            try:
                temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
            except:
                temp = 0
            
            # 전력
            try:
                power = pynvml.nvmlDeviceGetPowerUsage(handle) / 1000.0  # mW to W
            except:
                power = 0
            
            # 프로세스 정보
            processes = []
            if PSUTIL_AVAILABLE:
                try:
                    procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
                    for proc in procs:
                        try:
                            process = psutil.Process(proc.pid)
                            processes.append({
                                'pid': proc.pid,
                                'name': process.name(),
                                'memory': proc.usedGpuMemory / 1024**2,  # MB
                                'username': process.username()
                            })
                        except:
                            pass
                except:
                    pass
            
            gpus.append({
                'index': i,
                'name': name,
                'utilization': util.gpu,
                'memory_utilization': util.memory,
                'memory_used': mem_info.used / 1024**2,  # MB
                'memory_total': mem_info.total / 1024**2,  # MB
                'temperature': temp,
                'power': power,
                'processes': processes
            })
        
        pynvml.nvmlShutdown()
        return {'gpus': gpus}
    except Exception as e:
        print(f"Error getting pynvml data: {e}")
        return {}


def get_nvidia_smi_data() -> Dict[str, Any]:
    """nvidia-smi를 사용하여 GPU 정보를 수집합니다 (fallback)."""
    try:
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=index,name,utilization.gpu,utilization.memory,memory.used,memory.total,temperature.gpu,power.draw', '--format=csv,noheader,nounits'],
            capture_output=True,
            text=True,
            timeout=10
        )
        
        if result.returncode == 0:
            gpus = []
            for line in result.stdout.strip().split('\n'):
                parts = [p.strip() for p in line.split(',')]
                if len(parts) >= 8:
                    gpus.append({
                        'index': int(parts[0]),
                        'name': parts[1],
                        'utilization': float(parts[2]),
                        'memory_utilization': float(parts[3]),
                        'memory_used': float(parts[4]),
                        'memory_total': float(parts[5]),
                        'temperature': float(parts[6]),
                        'power': float(parts[7])
                    })
            return {'gpus': gpus}
    except Exception as e:
        print(f"Error getting nvidia-smi data: {e}")
    
    return {}


def send_data(server_url: str, server_name: str, data: Dict[str, Any]) -> bool:
    """수집한 데이터를 중앙 서버로 전송합니다."""
    try:
        payload = {
            'server_name': server_name,
            'timestamp': datetime.now().isoformat(),
            'hostname': os.uname().nodename,
            'gpu_data': data
        }
        
        response = requests.post(
            f"{server_url}/api/data",
            json=payload,
            timeout=10
        )
        
        if response.status_code == 200:
            return True
        else:
            print(f"Error sending data: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        print(f"Error sending data to server: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description='Nvitop Monitoring Agent')
    parser.add_argument('--server-url', required=True, help='중앙 서버 URL (예: http://localhost:5000)')
    parser.add_argument('--server-name', required=True, help='서버 이름')
    parser.add_argument('--interval', type=int, default=2, help='업데이트 간격 (초)')
    
    args = parser.parse_args()
    
    print(f"Starting agent for server: {args.server_name}")
    print(f"Connecting to: {args.server_url}")
    print(f"Update interval: {args.interval}s")
    
    while True:
        try:
            data = get_nvitop_data()
            if data:
                success = send_data(args.server_url, args.server_name, data)
                if success:
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] Data sent successfully")
                else:
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] Failed to send data")
            else:
                print(f"[{datetime.now().strftime('%H:%M:%S')}] No data collected")
            
            time.sleep(args.interval)
        except KeyboardInterrupt:
            print("\nShutting down agent...")
            break
        except Exception as e:
            print(f"Error in main loop: {e}")
            time.sleep(args.interval)


if __name__ == '__main__':
    main()
