from ruamel.yaml import YAML, comments
from ruamel.yaml.scalarstring import DoubleQuotedScalarString as dq
import os


base_dir = "/home/rpetrovich/code/Repos/lmru--devops--argocd-apps/"

yaml = YAML()
yaml.preserve_quotes = True
yaml.explicit_start = True
yaml.indent(mapping=2, sequence=4, offset=2)
yaml.allow_duplicate_keys = True

def load_yaml(file_path):
    with open(file_path, 'r') as file:
        data = list(yaml.load_all(file))
    return data

def save_yaml(data, file_path):
    with open(file_path, 'w') as file:
        yaml.dump_all(data, file)

def extract_cluster_info(node_path):
    parts = node_path.split('/')
    nodes_index = parts.index('nodes')
    env = parts[nodes_index - 2]
    cluster_index = parts.index('clusters') + 1
    cluster_name = parts[cluster_index]
    subpath = '-'.join(parts[cluster_index + 2:nodes_index - 1]) if nodes_index - cluster_index > 2 else ''
    return cluster_name, env, subpath

def add_label_if_needed(node_group):
    excluded_labels = [
        "d8-", "opendistro", "opensearch", "otel-", "system",
        "mimir-", "prometheus", "frontend", "elastic", "master",
        "loadbalancer", "deckhouse", "router", "moon-", "sentry",
        "odfe-", "obs-grafana", "devops-core-argocd", "victoria-",
        "clickhouse", "redis", "rabbit", "loki","cassandra", "jaeger", "vm-",
        "memcached", "kafka", "argocd", "keycloak", "ghe"
    ]
    
    name = node_group['metadata']['name']
    print(f"Processing NodeGroup: {name}")
    
    if any(excluded in name for excluded in excluded_labels):
        print(f"Excluded by name: {name}")
        return False  # Пропускаем обработку, если имя в списке исключений
    
    node_template = node_group['spec'].setdefault('nodeTemplate', comments.CommentedMap())
    labels = node_template.setdefault('labels', comments.CommentedMap())
    filtered_labels = {k: v for k, v in labels.items() if not k.startswith("threshold.")}
    
    if any(excluded in label for label in filtered_labels for excluded in excluded_labels):
        print(f"Excluded by existing labels: {filtered_labels}")
        return False  # Пропускаем обработку, если метки в списке исключений

    new_label_key = 'node-role.k8s.lmru.tech/application'
    
    # Добавляем логику: проверяем метку, но продолжаем проверку тейнтов
    label_already_present = new_label_key in labels
    if label_already_present:
        print(f"Label already present for {name}, but checking taints...")
    else:
        labels[new_label_key] = ""
        print(f"Label added for {name}: {new_label_key}")
    
    # Возвращаем True для продолжения обработки, даже если метка уже есть
    return True

def find_cluster_file(base_dir, cluster_name, env):
    target_directory = f"{base_dir}/common/aqua-kube-enforcer/{env}/values/"
    
    if not os.path.exists(target_directory):
        print(f"Directory does not exist: {target_directory}")
        return None

    for file in os.listdir(target_directory):
        if cluster_name in file and env in file and file.endswith('cluster.yaml') and '-elk-' not in file:
            return os.path.join(target_directory, file)

    print(f"No matching cluster file found that contains both '{cluster_name}' and '{env}' in {target_directory}")
    return None

def process_yaml_file(node_file):
    if '/elk/' in node_file:
        print(f"Skipping: {node_file} due to '/elk/'")
        return
    
    cluster_name, env, _ = extract_cluster_info(node_file)
    data = load_yaml(node_file)
    
    if not data:
        print(f"No data found in {node_file}")
        return
    
    changed = False
    tolerations = []
    
    for doc in data:
        if doc is None:
            print(f"Skipping invalid or empty YAML document in {node_file}")
            continue
        
        # Проверяем метки и продолжаем даже если метка уже есть
        if doc.get('kind') == "NodeGroup" and add_label_if_needed(doc):
            changed = True
            taints = doc.get('spec', {}).get('nodeTemplate', {}).get('taints', [])
            if taints:
                print(f"Taints found in NodeGroup {doc['metadata']['name']}: {taints}")
            for taint in taints:
                if 'key' in taint and 'value' in taint:
                    tolerations.append({
                        "key": taint['key'],
                        "operator": "Equal",
                        "effect": dq(taint.get('effect', '')),
                        "value": taint['value']
                    })
                    print(f"Adding toleration for {doc['metadata']['name']}: {taint}")
    
    if changed:
        save_yaml(data, node_file)
        print(f"Changes saved to {node_file}")
        
        if tolerations:
            cluster_file = find_cluster_file(base_dir, cluster_name, env)
            if cluster_file:
                print(f"Found cluster file for {cluster_name} in {env}: {cluster_file}")
                update_cluster_file(cluster_file, tolerations)
            else:
                print(f"No cluster file found for {cluster_name} in {env}")
        else:
            print(f"No tolerations found to add for {node_file}")


def update_cluster_file(cluster_file, new_tolerations):
    if not cluster_file:
        print("No matching cluster file found.")
        return
    
    print(f"Attempting to update cluster file: {cluster_file} with tolerations: {new_tolerations}")
    
    try:
        data = load_yaml(cluster_file)
        updated = False  # Флаг для отслеживания изменений
        
        for doc in data:
            # Убедимся, что раздел enforcer присутствует
            enforcer = doc.setdefault('enforcer', comments.CommentedMap())
            
            # Добавляем nodeSelector, если отсутствует
            enforcer.setdefault('nodeSelector', {})['node-role.k8s.lmru.tech/application'] = ''
            
            # Проверяем и добавляем новые tolerations
            existing_tolerations = enforcer.setdefault('tolerations', [])
            for new_tol in new_tolerations:
                # Проверяем, если такая же toleration уже присутствует
                if not any(
                    tol['key'] == new_tol['key'] and
                    tol['operator'] == new_tol['operator'] and
                    tol['effect'] == new_tol['effect'] and
                    tol['value'] == new_tol['value']
                    for tol in existing_tolerations
                ):
                    existing_tolerations.append(new_tol)
                    updated = True  # Устанавливаем флаг, что произошли изменения

        if updated:
            save_yaml(data, cluster_file)
            print(f"Updated cluster file successfully: {cluster_file}")
        else:
            print(f"No updates were needed for {cluster_file}")

    except Exception as e:
        print(f"Failed to update cluster file {cluster_file}: {e}")

def process_directory(directory):
    for root, _, files in os.walk(directory, topdown=True):
        # if "stage" not in root: # comment if you don't want to restrict "stage" only
        #     continue
        for file in files:
            if file == 'node_groups.yaml':
                full_path = os.path.join(root, file)
                print(f"Processing: {full_path}")
                process_yaml_file(full_path)

process_directory(base_dir)
