#!/usr/bin/env python
import yaml
import subprocess
import json
from getpass import getpass
import argparse


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Ansible-playbook wrapper for dynamic password input."
    )
    parser.add_argument("playbook", help="The path to the Ansible playbook.")
    parser.add_argument(
        "-i",
        "--inventory",
        default="inventory.yml",
        help="The path to the inventory file.",
    )
    return parser.parse_args()


def load_playbook_targets(playbook_path):
    # Load the playbook YAML and extract the hosts
    with open(playbook_path, "r") as f:
        playbook_data = yaml.safe_load(f)

    return playbook_data[0].get("hosts", [])


def fetch_inventory(inventory_path):
    inventory_json = subprocess.check_output(
        ["ansible-inventory", "-i", inventory_path, "--list"]
    )
    return json.loads(inventory_json)


def main():
    args = parse_arguments()
    playbook_targets = load_playbook_targets(args.playbook)
    inventory = fetch_inventory(args.inventory)

    # Determine targeted hosts
    targeted_hosts = set()

    if isinstance(playbook_targets, str):
        if playbook_targets in inventory:
            targeted_hosts.update(inventory[playbook_targets]["hosts"])
        elif playbook_targets == "all":
            targeted_hosts.update(inventory["_meta"]["hostvars"].keys())
    elif isinstance(playbook_targets, list):
        targeted_hosts.update(playbook_targets)

    # Validate targeted hosts against inventory
    valid_hosts = set(inventory["_meta"]["hostvars"].keys())
    targeted_hosts.intersection_update(valid_hosts)

    # Process each targeted host
    for host in targeted_hosts:
        print(f"Processing host: {host}")
        become_pass = getpass(f"Enter become password for {host}: ")

        # Execute Ansible playbook for each host
        subprocess.run(
            [
                "ansible-playbook",
                "-i",
                args.inventory,
                "--limit",
                host,
                args.playbook,
                "--extra-vars",
                f"ansible_become_pass={become_pass}",
            ]
        )


if __name__ == "__main__":
    main()