#!/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. This function is non-destructive and will only create or overwrite files for hosts defined in the cluster config."""
    # Create host_vars directory if it doesn't exist
    os.makedirs(host_vars_dir, exist_ok=True)
    
    # 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_path):
    """Generate group-level variables"""
    # Create parent directory if it doesn't exist
    all_vars_dir = os.path.dirname(group_vars_path)
    os.makedirs(all_vars_dir, exist_ok=True)

    # Remove the specific generated file if it exists to avoid stale data.
    if os.path.exists(group_vars_path):
        os.remove(group_vars_path)
    
    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']

    # Combine master and worker nodes to create a hostvars-like structure
    all_nodes = {**cluster_config.get('master', {}), **cluster_config.get('workers', {})}

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

    with open(group_vars_path, '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: ./tools/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)

    # Derive environment name from config filename (e.g., cluster.stress.yml -> stress)
    base_name = os.path.basename(config_path)
    if base_name == 'cluster.yml':
        env_name = ''
    elif base_name.startswith('cluster.') and base_name.endswith('.yml'):
        env_name = base_name[len('cluster.'):-len('.yml')]
    else:
        print(f"Warning: Unconventional config file name '{base_name}'. Using base name as environment identifier.")
        env_name = os.path.splitext(base_name)[0]

    # Define output paths based on environment
    inventory_suffix = f".{env_name}" if env_name else ""
    inventory_path = f"ansible/inventory{inventory_suffix}.ini"

    vars_suffix = f".{env_name}" if env_name else ""
    group_vars_path = f"ansible/group_vars/all/generated_vars{vars_suffix}.yml"
    
    # Load cluster configuration
    cluster_config = load_cluster_config(config_path)
    
    # Generate inventory file
    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
    generate_group_vars(cluster_config, group_vars_path)
    print(f"Generated group variables in {os.path.dirname(group_vars_path)}")

    print("Inventory generation complete!")

if __name__ == "__main__":
    main()