#!/usr/bin/env python3

import sys
import re
import subprocess
import json
import argparse
import csv

def run_local_cmd(cmd):
    try:
        result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=10)
        return result.stdout.strip() if result.returncode == 0 else ""
    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):
    return [line.split()[0] for line in output.strip().splitlines()[1:] if line.strip()]

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

def collect_node(host, user, use_ssh=False, key=None):
    rows = []

    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\"}'")

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

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

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

    return rows

def parse_config_line(cfg_lines, vmid, vmtype):
    def extract(field, default="-"):
        return next((line.split(":", 1)[1].strip() for line in cfg_lines if line.startswith(f"{field}:")), 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.startswith("tags:") for line in cfg_lines) else "No"

    storage_parts = [
        m.group(1)
        for line in cfg_lines
        if "none" not in line and any(d in line for d in ["scsi", "ide", "virtio", "sata", "mp", "rootfs"])
        if (m := re.search(r"size=([^, ]+)", line))
    ]
    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 output_table(rows, columns):
    if not rows:
        print("Keine Daten vorhanden.")
        return

    # Dynamische Breite pro Spalte
    col_widths = {
        col: max(len(col), max(len(str(row.get(col, '-'))) for row in rows))
        for col in columns
    }

    # Header
    header = " | ".join(col.ljust(col_widths[col]) for col in columns)
    separator = "-+-".join("-" * col_widths[col] for col in columns)
    print(header)
    print(separator)

    # Rows
    for row in rows:
        print(" | ".join(str(row.get(col, '-')).ljust(col_widths[col]) for col in columns))

def output_json(rows):
    print(json.dumps(rows, indent=2))

def output_csv(rows, columns):
    writer = csv.DictWriter(sys.stdout, fieldnames=columns)
    writer.writeheader()
    for row in rows:
        writer.writerow({col: row.get(col, '') for col in columns})

def main():
    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 Benutzer")
    parser.add_argument("--key", help="Pfad zum SSH Key")
    parser.add_argument("--columns", help="Welche Spalten anzeigen (Komma getrennt)")
    parser.add_argument("--output", choices=["cli", "json", "csv"], default="cli", help="Ausgabeformat")
    args = parser.parse_args()

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

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

    columns = args.columns.split(",") if args.columns else ['ID', 'Hostname', 'IP', 'Typ', 'CPU', 'RAM', 'Storage', 'Tags']

    if args.output == "cli":
        output_table(all_rows, columns)
    elif args.output == "json":
        output_json(all_rows)
    elif args.output == "csv":
        output_csv(all_rows, columns)

if __name__ == "__main__":
    main()