import yaml
import sys

def merge_templates(base_file, custom_file, output_file):
    print(f"Loading base template: {base_file}")
    with open(base_file, 'r', encoding='utf-8') as f:
        base_data = yaml.safe_load(f)

    print(f"Loading custom template: {custom_file}")
    with open(custom_file, 'r', encoding='utf-8') as f:
        custom_data = yaml.safe_load(f)

    # Assume standard Zabbix 6.0+ YAML export structure
    # zabbix_export -> templates -> [list]
    
    if 'zabbix_export' not in base_data or 'templates' not in base_data['zabbix_export']:
        print("Invalid base template structure.")
        return

    base_template = base_data['zabbix_export']['templates'][0]
    
    if 'zabbix_export' in custom_data and 'templates' in custom_data['zabbix_export']:
        custom_template = custom_data['zabbix_export']['templates'][0]
    else:
        print("Invalid custom template structure.")
        return

    print("Merging Discovery Rules...")
    if 'discovery_rules' in custom_template:
        if 'discovery_rules' not in base_template:
            base_template['discovery_rules'] = []
        
        # Check for duplicates by name or key to avoid collision
        existing_keys = {dr['key'] for dr in base_template['discovery_rules']}
        
        for dr in custom_template['discovery_rules']:
            if dr['key'] in existing_keys:
                print(f"  Skipping duplicate key: {dr['key']}")
            else:
                print(f"  Adding rule: {dr['name']}")
                base_template['discovery_rules'].append(dr)

    print("Merging Items...")
    if 'items' in custom_template:
        if 'items' not in base_template:
            base_template['items'] = []
            
        existing_item_keys = {item['key'] for item in base_template['items']}
        
        for item in custom_template['items']:
             if item['key'] in existing_item_keys:
                 print(f"  Skipping duplicate item key: {item['key']}")
             else:
                 print(f"  Adding item: {item['name']}")
                 base_template['items'].append(item)
    
    # Update Template Name
    base_template['name'] = "Nginx by Zabbix agent (Gold)"
    
    print(f"Writing merged template to {output_file}")
    with open(output_file, 'w', encoding='utf-8') as f:
        yaml.dump(base_data, f, sort_keys=False, allow_unicode=True)
    print("Merge Complete.")

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python merge_nginx_templates.py <base_yaml> <custom_yaml> <output_yaml>")
    else:
        merge_templates(sys.argv[1], sys.argv[2], sys.argv[3])