#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import concurrent.futures as futures
import os
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import Dict, Optional, Tuple, List

# ===== 설정: 기본 계정/비밀번호 (환경변수로 덮어쓰기 가능) =====
IDRAC_USER = os.getenv("IDRAC_USER", "root")
IDRAC_PASS = os.getenv("IDRAC_PASS", "calvin")

OUTPUT_DIR = Path("idrac_info")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

MAC_RE = re.compile(r"([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}")

# --- 유틸: 명령 실행 ---
def run(cmd: List[str]) -> Tuple[int, str]:
    try:
        p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, check=False)
        return p.returncode, p.stdout or ""
    except Exception as e:
        return 1, f"__EXEC_ERROR__ {e}"

# --- 파서: getsysinfo/swinventory/hwinventory 에서 필요한 값 추출 ---
def extract_first_mac(text: str) -> Optional[str]:
    m = MAC_RE.search(text)
    return m.group(0) if m else None

def parse_svc_tag(getsysinfo: str) -> Optional[str]:
    # 예: "SVC Tag               = ABCDEF2"
    m = re.search(r"SVC\s*Tag\s*=\s*([^\s]+)", getsysinfo)
    return m.group(1).strip() if m else None

def parse_specific_mac_in_getsysinfo(getsysinfo: str, key: str) -> Optional[str]:
    """
    key 예시:
      - "NIC.Integrated.1-1-1"
      - "NIC.Embedded.1-1-1"
    해당 라인에서 MAC 패턴 추출
    """
    for line in getsysinfo.splitlines():
        if key in line:
            mac = extract_first_mac(line)
            if mac:
                return mac
    return None

def parse_idrac_mac(getsysinfo: str) -> Optional[str]:
    """
    원본 스크립트는 "MAC Address             = " 를 grep 했습니다.
    해당 라인(들) 중 첫 MAC을 사용합니다.
    """
    lines = [ln for ln in getsysinfo.splitlines() if "MAC Address" in ln]
    for ln in lines:
        mac = extract_first_mac(ln)
        if mac:
            return mac
    return None

def parse_infiniband_slot_macs_from_swinventory(swinventory: str, slots: List[int]) -> Dict[int, Optional[str]]:
    """
    bash의 awk 트릭(해당 FQDD의 이전 줄에서 MAC 추출)을 그대로 재현:
    - swinventory 라인을 순회하면서 직전 라인을 prev로 저장
    - line이 'FQDD = InfiniBand.Slot.<N>-1' 에 매치되면 prev에서 MAC 추출
    """
    want = {s: None for s in slots}
    prev = ""
    for line in swinventory.splitlines():
        m = re.search(r"FQDD\s*=\s*InfiniBand\.Slot\.(\d+)-1", line)
        if m:
            slot = int(m.group(1))
            if slot in want and want[slot] is None:
                want[slot] = extract_first_mac(prev)  # 이전 줄에서 MAC
        prev = line
    return want

def parse_memory_vendor_initials(hwinventory: str) -> str:
    """
    원본: grep -A5 "DIMM" | grep "Manufacturer" | ... | sort | uniq | cut -c1
    - 간단화: DIMM 근처 5줄 윈도우에서 Manufacturer 라인 모으기
    - 제조사 첫 글자만 모아 중복 제거, 정렬 후 이어붙임
    """
    lines = hwinventory.splitlines()
    idxs = [i for i, ln in enumerate(lines) if "DIMM" in ln]
    vendors: List[str] = []
    for i in idxs:
        window = lines[i : i + 6]  # 본문 포함 6줄(= -A5)
        for ln in window:
            if "Manufacturer" in ln and "=" in ln:
                v = ln.split("=", 1)[1].strip()
                if v:
                    vendors.append(v)
    initials = sorted({v[0] for v in vendors if v})
    return "".join(initials)

def parse_ssd_manufacturers(hwinventory: str) -> str:
    """
    원본: grep -A3 "Disk.Bay" | grep "Manufacturer" | uniq
    - "Disk.Bay" 라인부터 3줄 윈도우 내 Manufacturer 라인 추출, 고유값 join
    """
    lines = hwinventory.splitlines()
    vendors: List[str] = []
    for i, ln in enumerate(lines):
        if "Disk.Bay" in ln:
            window = lines[i : i + 4]
            for w in window:
                if "Manufacturer" in w and "=" in w:
                    v = w.split("=", 1)[1].strip()
                    if v:
                        vendors.append(v)
    uniq = sorted({v for v in vendors})
    return ", ".join(uniq)

# --- 단일 IP 처리 ---
def fetch_idrac_info_for_ip(ip: str) -> Tuple[str, bool, str]:
    """
    returns: (ip, success, message)
    success=True면 파일 저장 완료
    """
    ip = ip.strip()
    if not ip:
        return ip, False, "빈 라인"

    # racadm 호출
    base = ["racadm", "-r", ip, "-u", IDRAC_USER, "-p", IDRAC_PASS]

    rc1, getsysinfo = run(base + ["getsysinfo"])
    rc2, swinventory = run(base + ["swinventory"])
    rc3, hwinventory = run(base + ["hwinventory"])

    if rc1 != 0 and rc2 != 0 and rc3 != 0:
        return ip, False, f"모든 racadm 호출 실패:\n{getsysinfo or ''}\n{swinventory or ''}\n{hwinventory or ''}"

    svc_tag = parse_svc_tag(getsysinfo) if getsysinfo else None
    if not svc_tag:
        return ip, False, "서비스 태그(SVC Tag) 추출 실패"

    # 개별 필드 파싱
    integrated_1 = parse_specific_mac_in_getsysinfo(getsysinfo, "NIC.Integrated.1-1-1") if getsysinfo else None
    integrated_2 = parse_specific_mac_in_getsysinfo(getsysinfo, "NIC.Integrated.1-2-1") if getsysinfo else None
    integrated_3 = parse_specific_mac_in_getsysinfo(getsysinfo, "NIC.Integrated.1-3-1") if getsysinfo else None
    integrated_4 = parse_specific_mac_in_getsysinfo(getsysinfo, "NIC.Integrated.1-4-1") if getsysinfo else None

    onboard_1 = parse_specific_mac_in_getsysinfo(getsysinfo, "NIC.Embedded.1-1-1") if getsysinfo else None
    onboard_2 = parse_specific_mac_in_getsysinfo(getsysinfo, "NIC.Embedded.2-1-1") if getsysinfo else None

    idrac_mac = parse_idrac_mac(getsysinfo) if getsysinfo else None

    # InfiniBand.Slot.<N>-1 MAC (이전 라인에서 MAC 추출) — 출력 순서 유지
    desired_slots = [38, 39, 37, 36, 32, 33, 34, 35, 31, 40]
    slot_macs = parse_infiniband_slot_macs_from_swinventory(swinventory, desired_slots) if swinventory else {s: None for s in desired_slots}

    memory_initials = parse_memory_vendor_initials(hwinventory) if hwinventory else ""
    ssd_vendors = parse_ssd_manufacturers(hwinventory) if hwinventory else ""

    # 저장
    out_path = OUTPUT_DIR / f"{svc_tag}.txt"
    with out_path.open("w", encoding="utf-8") as f:
        # 원본 스크립트와 동일한 출력 순서
        f.write(f"{svc_tag}\n")
        f.write(f"{integrated_1 or ''}\n")
        f.write(f"{integrated_2 or ''}\n")
        f.write(f"{integrated_3 or ''}\n")
        f.write(f"{integrated_4 or ''}\n")
        f.write(f"{onboard_1 or ''}\n")
        f.write(f"{onboard_2 or ''}\n")
        # 슬롯 고정 순서
        for sl in desired_slots:
            f.write(f"{slot_macs.get(sl) or ''}\n")
        f.write(f"{idrac_mac or ''}\n")
        f.write(f"{memory_initials}\n")
        f.write(f"{ssd_vendors}\n")

    return ip, True, f"저장: {out_path}"

# --- 메인 ---
def main():
    parser = argparse.ArgumentParser(description="iDRAC 정보 수집 (Python 포팅)")
    parser.add_argument("ip_file", help="IP 주소 목록 파일 (줄 단위)")
    parser.add_argument("--workers", type=int, default=20, help="병렬 스레드 수 (기본 20)")
    args = parser.parse_args()

    ip_file = Path(args.ip_file)
    if not ip_file.exists():
        print(f"IP 파일이 존재하지 않습니다: {ip_file}", file=sys.stderr)
        sys.exit(1)

    with ip_file.open("r", encoding="utf-8") as f:
        ips = [ln.strip() for ln in f if ln.strip()]

    if not ips:
        print("IP 목록이 비어 있습니다.", file=sys.stderr)
        sys.exit(1)

    print(f"[시작] 총 {len(ips)}대, workers={args.workers}, IDRAC_USER={IDRAC_USER}")

    t0 = time.time()
    ok = 0
    fail = 0

    # 병렬 수집
    with futures.ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(fetch_idrac_info_for_ip, ip): ip for ip in ips}
        for fut in futures.as_completed(futs):
            ip = futs[fut]
            try:
                _ip, success, msg = fut.result()
                prefix = "[OK] " if success else "[FAIL] "
                print(prefix + ip + " - " + msg)
                ok += int(success)
                fail += int(not success)
            except Exception as e:
                print(f"[EXC] {ip} - {e}", file=sys.stderr)
                fail += 1

    dt = int(time.time() - t0)
    h, r = divmod(dt, 3600)
    m, s = divmod(r, 60)

    print("\n정보 수집 완료.")
    print(f"성공 {ok}대 / 실패 {fail}대")
    print(f"수집 완료 시간: {h} 시간, {m} 분, {s} 초.")

if __name__ == "__main__":
    main()