#!/usr/bin/env python3
import argparse
import json
import sys

class DynamicInventory:
    def __init__(self):
        self.inventory = {
            "_meta": {"hostvars": {}},
            "all": {
                "children": ["ungrouped"],
                "vars": {"ansible_user": "admin", "datacenter": "us-east-1"}
            },
            "ungrouped": {"hosts": []}
        }
        self.roles = ["loadbalancer", "webserver", "database", "cache"]

    def _is_prime(self, n):
        if n <= 1: return False
        for i in range(2, int(n**0.5) + 1):
            if n % i == 0: return False
        return True

    def generate_data(self):
        for i in range(1, 13):
            hostname = f"srv-{i:03d}"
            
            role_idx = (i - 1) % len(self.roles)
            group_name = self.roles[role_idx]

            if group_name not in self.inventory:
                self.inventory[group_name] = {"hosts": [], "vars": {}}
                self.inventory["all"]["children"].append(group_name)

            self.inventory[group_name]["hosts"].append(hostname)

            # Assign specific group variables
            if group_name == "database":
                self.inventory[group_name]["vars"]["db_port"] = 5432
            elif group_name == "webserver":
                self.inventory[group_name]["vars"]["http_port"] = 80

            is_maintenance = self._is_prime(i)
            subnet = 10 + role_idx
            
            self.inventory["_meta"]["hostvars"][hostname] = {
                "ansible_host": f"192.168.{subnet}.{i}",
                "stage": "production" if i > 4 else "staging",
                "maintenance_mode": is_maintenance,
                "app_config": {
                    "threads": i * 2,
                    "max_connections": 1000 if group_name == "loadbalancer" else 200
                },
                "tags": ["legacy"] if i % 5 == 0 else ["modern", "cloud-native"]
            }

    def get_inventory(self):
        self.generate_data()
        return self.inventory

    def get_host(self, host):
        self.generate_data()
        return self.inventory["_meta"]["hostvars"].get(host, {})

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--list', action='store_true')
    parser.add_argument('--host', action='store')
    args = parser.parse_args()

    inv = DynamicInventory()

    if args.list:
        print(json.dumps(inv.get_inventory(), indent=2))
    elif args.host:
        print(json.dumps(inv.get_host(args.host), indent=2))
    else:
        sys.exit(1)

if __name__ == '__main__':
    main()