import yaml
import argparse
import sys
import re

def load_yaml(path):
    with open(path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

def save_yaml(data, path):
    with open(path, 'w', encoding='utf-8') as f:
        yaml.dump(data, f, sort_keys=False, indent=2, width=float("inf"), allow_unicode=True)

def normalize_expression(expression):
    """
    Removes the /Template Name/ part from item paths in the expression.
    Example: min(/Windows by Zabbix agent/system.cpu.util,5m) -> min(//system.cpu.util,5m)
    """
    # Regex to match /Template Name/Key
    # It usually looks like: /Host or Template/Key
    # We want to replace the first part with empty string if it starts with /
    if not expression:
        return ""
    
    # We use a simple heuristic: replace /[^/]+/ with //
    # But we must be careful not to break the key. 
    # Zabbix expression format: function(/Host/Key, params)
    # We want function(//Key, params)
    
    # Pattern: Look for / followed by anything not / followed by /
    return re.sub(r'\/[^\/]+\/', '//', expression)

def build_map(template_data):
    """
    Builds a map of entities from the template data.
    """
    mapping = {
        'items': {},
        'item_details': {}, # Store full item dict for property syncing
        'discovery_rules': {},
        'triggers': {},
        'item_prototypes': {}, # Keyed by DiscoveryRuleKey:ItemKey
        'trigger_prototypes': {}, # Keyed by DiscoveryRuleKey:Expression
        'template_uuid': None
    }
    
    if 'zabbix_export' not in template_data or 'templates' not in template_data['zabbix_export']:
        return mapping

    tmpl = template_data['zabbix_export']['templates'][0]
    mapping['template_uuid'] = tmpl.get('uuid')

    # Items
    for item in tmpl.get('items', []):
        mapping['items'][item['key']] = item.get('uuid')
        mapping['item_details'][item['key']] = item

    # Triggers
    for trigger in tmpl.get('triggers', []):
        norm_expr = normalize_expression(trigger['expression'])
        mapping['triggers'][norm_expr] = trigger.get('uuid')

    # Discovery Rules
    for rule in tmpl.get('discovery_rules', []):
        mapping['discovery_rules'][rule['key']] = rule.get('uuid')
        
        # Item Prototypes
        for proto in rule.get('item_prototypes', []):
            # Composite key: RuleKey:ProtoKey
            comp_key = f"{rule['key']}:{proto['key']}"
            mapping['item_prototypes'][comp_key] = proto.get('uuid')

            # Trigger Prototypes
            for trig_proto in proto.get('trigger_prototypes', []):
                 # Note: Trigger prototypes are nested under item prototypes in old formats, 
                 # but in 6.0+ they might be separate or nested.
                 # In the provided YAMLs, they are under item_prototypes for dependent items? 
                 # Wait, looking at file content...
                 # In file 7 (Active standard), trigger_prototypes are under discovery_rules 
                 # AND items can have trigger_prototypes (rare).
                 # Let's check the structure of file 7 again.
                 pass
        
        # Trigger Prototypes (Directly under Discovery Rule)
        for trig_proto in rule.get('trigger_prototypes', []):
             norm_expr = normalize_expression(trig_proto['expression'])
             comp_key = f"{rule['key']}:{norm_expr}"
             mapping['trigger_prototypes'][comp_key] = trig_proto.get('uuid')

    return mapping

def sync_uuids(source_path, target_path, output_path):
    print(f"Loading Source: {source_path}")
    source_data = load_yaml(source_path)
    print(f"Loading Target: {target_path}")
    target_data = load_yaml(target_path)

    source_map = build_map(source_data)
    
    if not source_map['template_uuid']:
        print("Error: No template found in source.")
        return

    # Sync Process
    target_tmpl = target_data['zabbix_export']['templates'][0]
    
    print(f"Syncing Template UUID: {source_map['template_uuid']}")
    target_tmpl['uuid'] = source_map['template_uuid']

    stats = {'items': 0, 'triggers': 0, 'discovery': 0, 'item_proto': 0, 'trigger_proto': 0}

    # Sync Items
    for item in target_tmpl.get('items', []):
        if item['key'] in source_map['items']:
            item['uuid'] = source_map['items'][item['key']]
            stats['items'] += 1
            
            # Additional Fix: Copy critical fields if missing in Target (Gold) 
            # but present in Source (Base). This fixes invalid/incomplete Gold items.
            source_item = source_map['item_details'].get(item['key'])
            if source_item:
                for field in ['type', 'delay', 'value_type', 'units', 'trends']:
                    if field not in item and field in source_item:
                        item[field] = source_item[field]
                        # Log if we are patching to debug
                        # print(f"Patched {field} for {item['key']}")

    # Sync Triggers
    for trigger in target_tmpl.get('triggers', []):
        norm_expr = normalize_expression(trigger['expression'])
        # Try exact match first (in case descriptions differ but expression is same)
        # Note: In source map we keyed by expression.
        if norm_expr in source_map['triggers']:
            trigger['uuid'] = source_map['triggers'][norm_expr]
            stats['triggers'] += 1

    # Sync Discovery Rules
    for rule in target_tmpl.get('discovery_rules', []):
        if rule['key'] in source_map['discovery_rules']:
            rule['uuid'] = source_map['discovery_rules'][rule['key']]
            stats['discovery'] += 1

            # Sync Item Prototypes
            for proto in rule.get('item_prototypes', []):
                comp_key = f"{rule['key']}:{proto['key']}"
                if comp_key in source_map['item_prototypes']:
                    proto['uuid'] = source_map['item_prototypes'][comp_key]
                    stats['item_proto'] += 1
            
            # Sync Trigger Prototypes
            for trig_proto in rule.get('trigger_prototypes', []):
                norm_expr = normalize_expression(trig_proto['expression'])
                comp_key = f"{rule['key']}:{norm_expr}"
                if comp_key in source_map['trigger_prototypes']:
                    trig_proto['uuid'] = source_map['trigger_prototypes'][comp_key]
                    stats['trigger_proto'] += 1

    print("Sync Complete.")
    print(f"Stats: {stats}")
    
    print(f"Saving to {output_path}")
    save_yaml(target_data, output_path)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Sync Zabbix Template UUIDs')
    parser.add_argument('--source', required=True, help='Path to Base Template (Source of UUIDs)')
    parser.add_argument('--target', required=True, help='Path to Gold Template (Target to update)')
    parser.add_argument('--output', required=True, help='Path to Output File')

    args = parser.parse_args()
    sync_uuids(args.source, args.target, args.output)