from config.basic_config import basic_config
from config.slurm_config import slurm_config
from config.shared_storage_config import shared_storage_config
from config.network_config import network_config
import os
import yaml


def remove_unassigned_configs(config):
    if isinstance(config, dict):
        return {k: remove_unassigned_configs(v) for k, v in config.items() if v is not None and remove_unassigned_configs(v) is not None}
    elif isinstance(config, list):
        return [remove_unassigned_configs(v) for v in config if v is not None and remove_unassigned_configs(v) is not None]
    else:
        return config


def config_menu():
    # Clear the screen
    os.system('cls' if os.name == 'nt' else 'clear')

    # Generate config for new cluster into a dictionary
    config = {}
    config.update(basic_config())
    config.update(slurm_config())
    config.update(shared_storage_config())
    subnet = network_config(config['Region'])

    config['HeadNode']['Networking'] = {"SubnetId": subnet}

    for s in config['Scheduling']['SlurmQueues']:
        s['Networking'] = {"SubnetIds": [subnet]}

    config = remove_unassigned_configs(config)

    with open('cluster-config.yaml', 'w') as f:
        yaml.safe_dump(config, f, default_flow_style=False)

    print('Configuration saved to cluster-config.yaml.')
