#!/usr/bin/env python3

import yaml
import sys
import os
import shutil
from jinja2 import Environment, FileSystemLoader

def load_cluster_config(config_path):
    """Load cluster configuration from YAML file"""
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def generate_inventory(cluster_config, inventory_path):
    """Generate Ansible inventory file from cluster configuration"""
    with open(inventory_path, 'w') as f:
        f.write("# This file is auto-generated by tools/generate-inventory.py\n")
        f.write("# Do not edit – your changes will be overwritten.\n")
        f.write("# Edit cluster.yml and re-run the generator instead.\n\n")
        
        # Master group
        f.write("[airflow_master]\n")
        for hostname, ip in cluster_config['master'].items():
            f.write(f"{hostname} ansible_host={ip}\n")
        
        f.write("\n")
        
        # Workers group
        f.write("[airflow_workers]\n")
        for hostname, config in cluster_config['workers'].items():
            f.write(f"{hostname} ansible_host={config['ip']}\n")

def generate_host_vars(cluster_config, host_vars_dir):
    """Generate host-specific variables"""
    # Create host_vars directory if it doesn't exist
    os.makedirs(host_vars_dir, exist_ok=True)
    
    # Clear existing host_vars files to avoid stale configurations
    for filename in os.listdir(host_vars_dir):
        file_path = os.path.join(host_vars_dir, filename)
        try:
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            print(f'Failed to delete {file_path}. Reason: {e}')
    
    # Get master IP for Redis configuration
    master_ip = list(cluster_config['master'].values())[0]
    
    # Generate master host variables
    for hostname, ip in cluster_config['master'].items():
        host_vars_file = os.path.join(host_vars_dir, f"{hostname}.yml")
        with open(host_vars_file, 'w') as f:
            f.write("---\n")
            f.write(f"# Variables for {hostname}\n")
            f.write(f"master_host_ip: {ip}\n")
            f.write(f"redis_port: 52909\n")

    # Generate worker proxy variables
    for hostname, config in cluster_config['workers'].items():
        host_vars_file = os.path.join(host_vars_dir, f"{hostname}.yml")
        worker_proxies = config.get('proxies', [])
        
        with open(host_vars_file, 'w') as f:
            f.write("---\n")
            f.write(f"# Variables for {hostname}\n")
            f.write(f"master_host_ip: {master_ip}\n")
            f.write("redis_port: 52909\n")
            if worker_proxies:
                f.write("worker_proxies:\n")
                for proxy in worker_proxies:
                    f.write(f"  - \"{proxy}\"\n")

def generate_group_vars(cluster_config, group_vars_dir):
    """Generate group-level variables"""
    # Create group_vars directory if it doesn't exist
    os.makedirs(group_vars_dir, exist_ok=True)
    
    # Create group_vars/all directory if it doesn't exist
    all_vars_dir = os.path.join(group_vars_dir, "all")
    os.makedirs(all_vars_dir, exist_ok=True)

    # Define path for the generated file and remove it if it exists to avoid stale data.
    # This is safer than removing the whole directory, which would delete vault.yml.
    all_vars_file = os.path.join(all_vars_dir, "generated_vars.yml")
    if os.path.exists(all_vars_file):
        os.remove(all_vars_file)
    external_ips = cluster_config.get('external_access_ips', [])

    # Get master IP for Redis configuration
    master_ip = list(cluster_config['master'].values())[0]

    with open(all_vars_file, 'w') as f:
        f.write("---\n")
        f.write("# This file is auto-generated by tools/generate-inventory.py\n")
        f.write("# Do not edit – your changes will be overwritten.\n")
        f.write(f"master_host_ip: {master_ip}\n")
        f.write("redis_port: 52909\n")
        f.write("external_access_ips:\n")
        if external_ips:
            for ip in external_ips:
                f.write(f"  - \"{ip}\"\n")
        else:
            f.write("  []\n")

def main():
    if len(sys.argv) != 2:
        print("Usage: python3 generate-inventory.py <cluster-config-file>")
        sys.exit(1)
    
    config_path = sys.argv[1]
    
    # Check if config file exists
    if not os.path.exists(config_path):
        print(f"Error: Configuration file {config_path} not found")
        sys.exit(1)
    
    # Load cluster configuration
    cluster_config = load_cluster_config(config_path)
    
    # Generate inventory file
    inventory_path = "ansible/inventory.ini"
    generate_inventory(cluster_config, inventory_path)
    print(f"Generated {inventory_path}")
    
    # Generate host variables
    host_vars_dir = "ansible/host_vars"
    generate_host_vars(cluster_config, host_vars_dir)
    print(f"Generated host variables in {host_vars_dir}")

    # Generate group variables
    group_vars_dir = "ansible/group_vars"
    generate_group_vars(cluster_config, group_vars_dir)
    print(f"Generated group variables in {group_vars_dir}")

    print("Inventory generation complete!")

if __name__ == "__main__":
    main()