import wmi
import requests
import json
import socket
import platform
import sys
import time

def collect_specs():
    try:
        c = wmi.WMI()
        computer = c.Win32_ComputerSystem()[0]
        os_info = c.Win32_OperatingSystem()[0]
        proc = c.Win32_Processor()[0]
        board = c.Win32_BaseBoard()[0]
        
        # 1. 상세 GPU 정보 수집 (모든 그래픽 카드)
        gpu_list = []
        for g in c.Win32_VideoController():
            gpu_list.append(g.Name)
        gpu_info = ", ".join(gpu_list) if gpu_list else "N/A"

        # 2. 모든 저장장치 정보 수집 및 SSD/HDD 구분
        storage_list = []
        
        # Windows 8 이상에서 작동하는 상세 저장소 정보 수집 시도
        physical_disks = {}
        try:
            storage_c = wmi.WMI(namespace="Root\\Microsoft\\Windows\\Storage")
            for d in storage_c.MSFT_PhysicalDisk():
                # MediaType: 3(HDD), 4(SSD), 0(Unspecified)
                physical_disks[d.DeviceId] = d.MediaType
        except:
            pass

        for d in c.Win32_DiskDrive():
            size_gb = round(float(d.Size) / (1024**3)) if d.Size else 0
            
            # 미디어 타입 판단
            media_type = physical_disks.get(str(d.Index), 0)
            prefix = ""
            if media_type == 4:
                prefix = "[SSD] "
            elif media_type == 3:
                prefix = "[HDD] "
            else:
                # 힌트가 없을 경우 모델명으로 추측
                cap = d.Caption.upper()
                if "SSD" in cap or "NVME" in cap or "FLASH" in cap:
                    prefix = "[SSD] "
                else:
                    prefix = "[HDD] "
            
            storage_list.append(f"{prefix}{d.Caption} ({size_gb}GB)")
        
        # DB 필드(SSD1, SSD2, SSD3)에 나눠 담기
        storage1 = storage_list[0] if len(storage_list) > 0 else "N/A"
        storage2 = storage_list[1] if len(storage_list) > 1 else ""
        storage3 = storage_list[2] if len(storage_list) > 2 else ""

        # 실시간 데이터 추출
        specs = {
            "메인보드": f"{board.Manufacturer} {board.Product}".strip(),
            "CPU": proc.Name.strip(),
            "RAM": f"{round(float(computer.TotalPhysicalMemory) / (1024**3))}GB",
            "OS": os_info.Caption,
            "GPU": gpu_info,
            "SSD1": storage1,
            "SSD2": storage2,
            "SSD3": storage3,
            "비고": "실시간 에이전트(EXE) 자동 수집"
        }
        return specs
    except Exception as e:
        print(f"데이터 수집 중 오류 발생: {e}")
        return None

def send_data(specs, server_url, asset_code):
    try:
        # 전송 데이터에 자산코드 추가 (식별용)
        specs["자산코드"] = asset_code
        print(f"\n📡 서버로 전송 중... ({server_url})")
        response = requests.post(server_url, json=specs, timeout=10)
        if response.status_code == 200:
            print("✅ 전송 성공! ITAM 시스템에서 확인하세요.")
        else:
            print(f"❌ 전송 실패: 서버 응답 코드 {response.status_code}")
    except Exception as e:
        print(f"❌ 서버 연결 오류: {e}")
        print("서버가 켜져 있는지, URL이 맞는지 확인해주세요.")

if __name__ == "__main__":
    print("========================================")
    print("   ITAM PC 실시간 사양 수집 에이전트 (v1.1)")
    print("========================================\n")
    
    # 1. 정보 수집
    print("🔍 하드웨어 정보를 읽어오는 중...")
    data = collect_specs()
    
    if data:
        print("\n[수집된 실제 사양]")
        display_map = {
            "메인보드": "메인보드",
            "CPU": "CPU",
            "RAM": "RAM",
            "OS": "OS",
            "GPU": "GPU",
            "SSD1": "Storage 1",
            "SSD2": "Storage 2",
            "SSD3": "Storage 3",
            "비고": "비고"
        }
        for key, value in data.items():
            if value: # 값이 있는 경우만 표시
                label = display_map.get(key, key)
                print(f" - {label}: {value}")
        
        print("\n" + "="*40)
        asset_code = input("등록할 자산번호를 입력하세요 (예: PC-001): ").strip()
        if not asset_code:
            print("❌ 자산번호 없이는 전송할 수 없습니다.")
        else:
            server_ip = input("서버 IP를 입력하세요 (기본값 localhost): ").strip()
            if not server_ip: server_ip = "localhost"
            
            target_url = f"http://{server_ip}:3000/api/agent/collect"
            
            confirm = input("\n위 정보를 서버로 전송할까요? (y/n): ")
            if confirm.lower() == 'y':
                send_data(data, target_url, asset_code)
    
    print("\n5초 후 프로그램이 종료됩니다...")
    time.sleep(5)