#!/usr/bin/env python3

import sys
import re
import subprocess
import json
import socket
import os

def run_local_cmd(cmd):
    try:
        result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=10)
        if result.returncode != 0:
            return ""
        return result.stdout.strip()
    except subprocess.TimeoutExpired:
        return ""

def run_ssh_cmd(host, user, cmd, key=None):
    ssh_cmd = f"ssh -o BatchMode=yes -o ConnectTimeout=5 "
    if key:
        ssh_cmd += f"-i {key} "
    ssh_cmd += f"{user}@{host} '{cmd}'"
    return run_local_cmd(ssh_cmd)

def parse_vm_list(output):
    vms = []
    for line in output.strip().splitlines()[1:]:
        fields = line.split()
        if len(fields) >= 2:
            vms.append(fields[0])
    return vms

def get_config_lines(host, user, use_ssh=False, key=None, vmid=None, ct=False):
    cmd = f"{'pct' if ct else 'qm'} config {vmid}"
    if use_ssh:
        return run_ssh_cmd(host, user, cmd, key).splitlines()
    return run_local_cmd(cmd).splitlines()

def collect_node(host, user, use_ssh=False, key=None):
    rows = []
    # Node Info
    hostname = run_ssh_cmd(host, user, "hostname -f", key) if use_ssh else run_local_cmd("hostname -f")
    ip = run_ssh_cmd(host, user, "hostname -I | awk '{print $1}'", key) if use_ssh else run_local_cmd("hostname -I | awk '{print $1}'")
    cpu = run_ssh_cmd(host, user, "nproc", key) if use_ssh else run_local_cmd("nproc")
    ram = run_ssh_cmd(host, user, "free -m | awk '/Mem:/ {print $2\"M\"}'", key) if use_ssh else run_local_cmd("free -m | awk '/Mem:/ {print $2\"M\"}'")

    node_info = {
        'ID': '0',
        'Hostname': hostname if hostname else '-',
        'IP': ip if ip else '-',
        'Typ': 'Node',
        'CPU': cpu if cpu else 'UNL',
        'RAM': ram if ram else '-',
        'Storage': "-",
        'Tags': "-"
    }
    rows.append(node_info)

    # VMs
    cmd_qm = "qm list"
    vmids = parse_vm_list(run_ssh_cmd(host, user, cmd_qm, key) if use_ssh else run_local_cmd(cmd_qm))
    for vmid in vmids:
        cfg = get_config_lines(host, user, use_ssh, key, vmid)
        row = parse_config_line(cfg, vmid, "VM")
        rows.append(row)

    # CTs
    cmd_pct = "pct list"
    ctids = parse_vm_list(run_ssh_cmd(host, user, cmd_pct, key) if use_ssh else run_local_cmd(cmd_pct))
    for ctid in ctids:
        cfg = get_config_lines(host, user, use_ssh, key, ctid, ct=True)
        row = parse_config_line(cfg, ctid, "CT")
        rows.append(row)

    return rows

def parse_config_line(cfg_lines, vmid, vmtype):
    cfg = "\n".join(cfg_lines)
    def extract(field, default="-"):
        for line in cfg_lines:
            if line.strip().startswith(f"{field}:"):
                return line.split(":", 1)[1].strip()
        return default

    name = extract("name", f"{vmtype}-{vmid}")
    hostname = extract("hostname", name)
    cores = extract("cores", "UNL")
    memory = extract("memory", "0") + "M"

    ip = "-"
    for line in cfg_lines:
        if "ip=" in line:
            ip = line.split("ip=")[1].split(',')[0]
            break

    tags = "Yes" if any(line.strip().startswith("tags:") for line in cfg_lines) else "No"

    storage_parts = []
    for line in cfg_lines:
        if any(disk in line for disk in ["scsi", "ide", "virtio", "sata", "mp", "rootfs"]):
            if "none" in line:
                continue
            m = re.search(r"size=([^, ]+)", line)
            if m:
                storage_parts.append(m.group(1))
    storage = "+".join(storage_parts) if storage_parts else "-"

    return {
        'ID': vmid,
        'Hostname': hostname,
        'IP': ip,
        'Typ': vmtype,
        'CPU': cores,
        'RAM': memory,
        'Storage': storage,
        'Tags': tags
    }

def print_table(rows):
    columns = ['ID', 'Hostname', 'IP', 'Typ', 'CPU', 'RAM', 'Storage', 'Tags']
    max_width = {
        'ID': 8,
        'Hostname': 25,
        'IP': 17,
        'Typ': 7,
        'CPU': 6,
        'RAM': 8,
        'Storage': 20,
        'Tags': 5,
    }

    # Header
    header = ""
    for col in columns:
        header += f"{col.ljust(max_width[col])} "
    print(header.strip())
    print('-' * (sum(max_width.values()) + len(columns) - 1))

    # Rows
    for row in rows:
        line = ""
        for col in columns:
            val = str(row.get(col, '-'))
            if len(val) > max_width[col]:
                val = val[:max_width[col]-3] + "..."
            line += f"{val.ljust(max_width[col])} "
        print(line.strip())

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Proxmox Inventory")
    parser.add_argument("--local", action="store_true", help="Nur lokale Node abfragen")
    parser.add_argument("--nodes", help="Komma-separierte Liste von Nodes")
    parser.add_argument("--user", default="root", help="SSH User")
    parser.add_argument("--key", help="Pfad zum SSH-Key")
    args = parser.parse_args()

    hosts = ["localhost"] if args.local else (args.nodes.split(",") if args.nodes else [])
    if not hosts:
        print("Fehler: Entweder --local oder --nodes muss angegeben werden.")
        sys.exit(1)

    use_ssh = not args.local

    all_rows = []
    for h in hosts:
        all_rows.extend(collect_node(h, args.user, use_ssh, args.key))

    print_table(all_rows)

if __name__ == "__main__":
    main()