#!/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, config in cluster_config['master'].items():
            line = f"{hostname} ansible_host={config['ip']}"
            if 'port' in config:
                line += f" ansible_port={config['port']}"
            f.write(line + "\n")
        
        f.write("\n")
        
        # Workers group
        f.write("[airflow_workers]\n")
        for hostname, config in cluster_config['workers'].items():
            line = f"{hostname} ansible_host={config['ip']}"
            if 'port' in config:
                line += f" ansible_port={config['port']}"
            f.write(line + "\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 from the new structure
    master_ip = list(cluster_config['master'].values())[0]['ip']
    
    # Get global proxy definitions
    shadowsocks_proxies = cluster_config.get('shadowsocks_proxies', {})

    # Combine master and worker nodes for processing
    all_nodes = {**cluster_config['master'], **cluster_config['workers']}

    for hostname, config in all_nodes.items():
        host_vars_file = os.path.join(host_vars_dir, f"{hostname}.yml")
        
        # Per-node list of proxies to USE
        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")
            
            # Write the global proxy definitions for deployment
            if shadowsocks_proxies:
                f.write("shadowsocks_proxies:\n")
                for name, proxy_config in shadowsocks_proxies.items():
                    f.write(f"  {name}:\n")
                    f.write(f"    server: \"{proxy_config['server']}\"\n")
                    f.write(f"    server_port: {proxy_config['server_port']}\n")
                    f.write(f"    local_port: {proxy_config['local_port']}\n")
                    f.write(f"    vault_password_key: \"{proxy_config['vault_password_key']}\"\n")

            # Write the per-node list of proxies to USE
            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)
    
    global_vars = cluster_config.get('global_vars', {})
    external_ips = cluster_config.get('external_access_ips', [])

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

    # Prepare data for YAML dump
    generated_data = {
        'master_host_ip': master_ip,
        'redis_port': 52909,
        'external_access_ips': external_ips if external_ips else []
    }
    generated_data.update(global_vars)

    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")
        yaml.dump(generated_data, f, default_flow_style=False)

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()