#!/usr/bin/env python3
import json
import sys
import os
import yaml
from pathlib import Path

def generate_inventory(discovery_data):
    hostname = discovery_data.get("hostname", "unknown-node")
    host_dir = Path(f"hosts/{hostname}")
    host_dir.mkdir(parents=True, exist_ok=True)

    # 1. host.yaml
    host_yaml = {
        "hostname": hostname,
        "roles": ["unassigned"],
        "network": {
            "tailscale_ip": discovery_data["network"]["tailscale"]["ip"]
        },
        "runtime": {
            "root": "/opt/homelab"
        },
        "deployment": {
            "mode": "pull",
            "managed_by": "saturn"
        }
    }
    with open(host_dir / "host.yaml", "w") as f:
        yaml.dump(host_yaml, f, sort_keys=False)

    # 2. capabilities.yaml
    capabilities_yaml = {
        "capabilities": {
            "hardware": {
                "cpu": {
                    "arch": discovery_data["os"]["arch"],
                    "cores": discovery_data["hardware"]["cpu"]["cores"],
                    "threads": discovery_data["hardware"]["cpu"]["threads"]
                },
                "memory": {
                    "total_gb": discovery_data["hardware"]["memory"]["total_gb"]
                },
                "acceleration": {
                    "type": "gpu" if discovery_data["hardware"]["gpu"]["present"] else "none"
                }
            },
            "virtualization": {
                "supported": discovery_data["virtualization"]["supported"],
                "type": discovery_data["virtualization"]["type"]
            },
            "storage": {
                "persistence": "persistent",
                "type": "ssd", # Default assumption
                "capacity_gb": sum([float(d["size"].rstrip("G")) for d in discovery_data["hardware"]["disks"] if "G" in d["size"]]) # Very rough estimate
            },
            "networking": {
                "reachability": "tailscale-only" if discovery_data["network"]["tailscale"]["status"] == "active" else "direct",
                "ingress_suitability": False,
                "bandwidth": "unknown"
            },
            "runtime": {
                "container_engine": "docker" if discovery_data["docker"]["available"] else "none",
                "os": discovery_data["os"]["distro"]
            }
        }
    }
    with open(host_dir / "capabilities.yaml", "w") as f:
        yaml.dump(capabilities_yaml, f, sort_keys=False)

    # 3. paths.yaml
    paths_yaml = {
        "host": hostname,
        "runtime_root": "/opt/homelab",
        "conventions": {
            "services": "/opt/homelab/services",
            "data": "/opt/homelab/data",
            "config": "/opt/homelab/config",
            "logs": "/opt/homelab/logs"
        }
    }
    with open(host_dir / "paths.yaml", "w") as f:
        yaml.dump(paths_yaml, f, sort_keys=False)

    # 4. networking.yaml
    networking_yaml = {
        "host": hostname,
        "uplink": {
            "type": "unknown",
            "connectivity": "unknown"
        },
        "tailscale": {
            "enabled": True if discovery_data["network"]["tailscale"]["status"] == "active" else False,
            "host_ip": discovery_data["network"]["tailscale"]["ip"],
            "role": "internal-management"
        }
    }
    with open(host_dir / "networking.yaml", "w") as f:
        yaml.dump(networking_yaml, f, sort_keys=False)

    print(f"Inventory generated for {hostname} in {host_dir}")

def main():
    if len(sys.argv) > 1:
        with open(sys.argv[1], "r") as f:
            data = json.load(f)
    else:
        # Read from stdin
        data = json.load(sys.stdin)
    
    generate_inventory(data)

if __name__ == "__main__":
    main()