#!/usr/bin/env python3

import json
import subprocess
import sys
from pathlib import Path

from jinja2 import Environment, FileSystemLoader


def get_terraform_outputs():
    result = subprocess.run(
        ["terraform", "output", "-json"],
        cwd="../terraform",
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        print(f"Error running terraform output: {result.stderr}")
        sys.exit(1)

    return json.loads(result.stdout)


def main():
    outputs = get_terraform_outputs()

    control_plane_names = outputs["control_plane_names"]["value"]
    control_plane_ips = outputs["control_plane_ips"]["value"]
    control_plane_private_ips = outputs["control_plane_private_ips"]["value"]
    worker_names = outputs["worker_names"]["value"]
    worker_ips = outputs["worker_ips"]["value"]
    worker_private_ips = outputs["worker_private_ips"]["value"]
    kube_api_lb_ip = outputs.get("kube_api_lb_ip", {}).get("value", control_plane_ips[0])

    control_planes = [
        {
            "name": name,
            "public_ip": public_ip,
            "private_ip": private_ip,
        }
        for name, public_ip, private_ip in zip(
            control_plane_names, control_plane_ips, control_plane_private_ips
        )
    ]

    workers = [
        {
            "name": name,
            "public_ip": public_ip,
            "private_ip": private_ip,
        }
        for name, public_ip, private_ip in zip(
            worker_names, worker_ips, worker_private_ips
        )
    ]

    data = {
        "control_planes": control_planes,
        "workers": workers,
        "private_key_file": outputs["ssh_private_key_path"]["value"],
        "kube_api_lb_ip": kube_api_lb_ip,
    }

    env = Environment(loader=FileSystemLoader("."))
    template = env.get_template("inventory.tmpl")
    inventory = template.render(**data)

    Path("inventory.ini").write_text(inventory)
    print("Generated inventory.ini")


if __name__ == "__main__":
    main()