#!/usr/bin/env python3
"""
Dynamic Ansible inventory script that extracts host information from the NixOS flake.

Generates groups:
  - flake_hosts: All hosts defined in the flake
  - tier_test, tier_prod: Hosts by deployment tier
  - role_<name>: Hosts by role (dns, vault, monitoring, etc.)

Usage:
  ./dynamic_flake.py --list    # Return full inventory
  ./dynamic_flake.py --host X  # Return host vars (not used, but required by Ansible)
"""

import json
import subprocess
import sys
from pathlib import Path


def get_flake_dir() -> Path:
    """Find the flake root directory."""
    script_dir = Path(__file__).resolve().parent
    # ansible/inventory/dynamic_flake.py -> repo root
    return script_dir.parent.parent


def evaluate_flake() -> dict:
    """Evaluate the flake and extract host metadata."""
    flake_dir = get_flake_dir()

    # Nix expression to extract relevant config from each host
    nix_expr = """
    configs: builtins.mapAttrs (name: cfg: {
      hostname = cfg.config.networking.hostName;
      domain = cfg.config.networking.domain or "home.2rjus.net";
      tier = cfg.config.homelab.host.tier;
      role = cfg.config.homelab.host.role;
      labels = cfg.config.homelab.host.labels;
      dns_enabled = cfg.config.homelab.dns.enable;
    }) configs
    """

    try:
        result = subprocess.run(
            [
                "nix",
                "eval",
                "--json",
                f"{flake_dir}#nixosConfigurations",
                "--apply",
                nix_expr,
            ],
            capture_output=True,
            text=True,
            check=True,
            cwd=flake_dir,
        )
        return json.loads(result.stdout)
    except subprocess.CalledProcessError as e:
        print(f"Error evaluating flake: {e.stderr}", file=sys.stderr)
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"Error parsing nix output: {e}", file=sys.stderr)
        sys.exit(1)


def sanitize_group_name(name: str) -> str:
    """Sanitize a string for use as an Ansible group name.

    Ansible group names should contain only alphanumeric characters and underscores.
    """
    return name.replace("-", "_")


def build_inventory(hosts_data: dict) -> dict:
    """Build Ansible inventory structure from host data."""
    inventory = {
        "_meta": {"hostvars": {}},
        "flake_hosts": {"hosts": []},
    }

    # Track groups we need to create
    tier_groups: dict[str, list[str]] = {}
    role_groups: dict[str, list[str]] = {}

    for _config_name, host_info in hosts_data.items():
        hostname = host_info["hostname"]
        domain = host_info["domain"]
        tier = host_info["tier"]
        role = host_info["role"]
        labels = host_info["labels"]
        dns_enabled = host_info["dns_enabled"]

        # Skip hosts that have DNS disabled (like templates)
        if not dns_enabled:
            continue

        # Skip hosts with ansible = "false" label
        if labels.get("ansible") == "false":
            continue

        fqdn = f"{hostname}.{domain}"

        # Use short hostname as inventory name, FQDN for connection
        inventory_name = hostname

        # Add to flake_hosts group
        inventory["flake_hosts"]["hosts"].append(inventory_name)

        # Add host variables
        inventory["_meta"]["hostvars"][inventory_name] = {
            "ansible_host": fqdn,  # Connect using FQDN
            "fqdn": fqdn,
            "tier": tier,
            "role": role,
        }

        # Group by tier
        tier_group = f"tier_{sanitize_group_name(tier)}"
        if tier_group not in tier_groups:
            tier_groups[tier_group] = []
        tier_groups[tier_group].append(inventory_name)

        # Group by role (if set)
        if role:
            role_group = f"role_{sanitize_group_name(role)}"
            if role_group not in role_groups:
                role_groups[role_group] = []
            role_groups[role_group].append(inventory_name)

    # Add tier groups to inventory
    for group_name, hosts in tier_groups.items():
        inventory[group_name] = {"hosts": hosts}

    # Add role groups to inventory
    for group_name, hosts in role_groups.items():
        inventory[group_name] = {"hosts": hosts}

    return inventory


def main():
    if len(sys.argv) < 2:
        print("Usage: dynamic_flake.py --list | --host <hostname>", file=sys.stderr)
        sys.exit(1)

    if sys.argv[1] == "--list":
        hosts_data = evaluate_flake()
        inventory = build_inventory(hosts_data)
        print(json.dumps(inventory, indent=2))
    elif sys.argv[1] == "--host":
        # Ansible calls this to get vars for a specific host
        # We provide all vars in _meta.hostvars, so just return empty
        print(json.dumps({}))
    else:
        print(f"Unknown option: {sys.argv[1]}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()