#!/usr/bin/env python3
"""
Enhance guest agent verification in all VM YAML templates.
Adds detailed verification commands matching the check script.
"""

import re
import sys
import os
from pathlib import Path
from datetime import datetime

# Enhanced verification block
ENHANCED_VERIFICATION = '''        # Verify packages are installed
        - |
          echo "=========================================="
          echo "Verifying required packages are installed..."
          echo "=========================================="
          for pkg in qemu-guest-agent curl wget net-tools chrony unattended-upgrades; do
            if ! dpkg -l | grep -q "^ii.*$pkg"; then
              echo "ERROR: Package $pkg is not installed"
              exit 1
            fi
            echo "✅ Package $pkg is installed"
          done
          echo "All required packages verified"
        
        # Verify qemu-guest-agent package details
        - |
          echo "=========================================="
          echo "Checking qemu-guest-agent package details..."
          echo "=========================================="
          if dpkg -l | grep -q "^ii.*qemu-guest-agent"; then
            echo "✅ qemu-guest-agent package IS installed"
            dpkg -l | grep qemu-guest-agent
          else
            echo "❌ qemu-guest-agent package is NOT installed"
            echo "Attempting to install..."
            apt-get update
            apt-get install -y qemu-guest-agent
          fi
        
        # Enable and start QEMU Guest Agent
        - |
          echo "=========================================="
          echo "Enabling and starting QEMU Guest Agent..."
          echo "=========================================="
          systemctl enable qemu-guest-agent
          systemctl start qemu-guest-agent
          echo "QEMU Guest Agent enabled and started"
        
        # Verify guest agent service is running
        - |
          echo "=========================================="
          echo "Verifying QEMU Guest Agent service status..."
          echo "=========================================="
          for i in {1..30}; do
            if systemctl is-active --quiet qemu-guest-agent; then
              echo "✅ QEMU Guest Agent service IS running"
              systemctl status qemu-guest-agent --no-pager -l
              exit 0
            fi
            echo "Waiting for QEMU Guest Agent to start... ($i/30)"
            sleep 1
          done
          echo "⚠️  WARNING: QEMU Guest Agent may not have started properly"
          systemctl status qemu-guest-agent --no-pager -l || true
          echo "Attempting to restart..."
          systemctl restart qemu-guest-agent
          sleep 3
          if systemctl is-active --quiet qemu-guest-agent; then
            echo "✅ QEMU Guest Agent started after restart"
          else
            echo "❌ QEMU Guest Agent failed to start"
          fi'''


def find_verification_block(content):
    """Find the old verification block in the content."""
    # Pattern to match from "Verify packages" to end of guest agent verification
    pattern = r'(        # Verify packages are installed.*?          systemctl status qemu-guest-agent --no-pager \|\| true)'
    match = re.search(pattern, content, re.DOTALL)
    return match


def enhance_file(file_path):
    """Enhance a single YAML file with improved verification."""
    print(f"📝 Processing: {file_path}")
    
    # Read file
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        print(f"❌ Error reading {file_path}: {e}")
        return False
    
    # Check if file contains guest agent verification
    if "Verifying required packages are installed" not in content:
        print(f"⏭️  Skipping {file_path} (no guest agent verification found)")
        return None
    
    # Check if already enhanced
    if "Checking qemu-guest-agent package details" in content:
        print(f"✅ Already enhanced: {file_path}")
        return None
    
    # Find and replace
    match = find_verification_block(content)
    if not match:
        print(f"⚠️  Could not find verification block in {file_path}")
        return False
    
    # Create backup
    backup_path = f"{file_path}.backup-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    try:
        with open(backup_path, 'w', encoding='utf-8') as f:
            f.write(content)
    except Exception as e:
        print(f"❌ Error creating backup: {e}")
        return False
    
    # Replace
    new_content = content[:match.start()] + ENHANCED_VERIFICATION + content[match.end():]
    
    # Write updated content
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(new_content)
        print(f"✅ Updated: {file_path}")
        return True
    except Exception as e:
        print(f"❌ Error writing {file_path}: {e}")
        # Restore from backup
        try:
            with open(backup_path, 'r', encoding='utf-8') as f:
                with open(file_path, 'w', encoding='utf-8') as out:
                    out.write(f.read())
        except:
            pass
        return False


def main():
    """Main function."""
    script_dir = Path(__file__).parent
    project_root = script_dir.parent
    templates_dir = project_root / "examples" / "production"
    
    if not templates_dir.exists():
        print(f"❌ Templates directory not found: {templates_dir}")
        sys.exit(1)
    
    print("==========================================")
    print("Enhancing Guest Agent Verification")
    print("==========================================")
    print(f"Target directory: {templates_dir}")
    print()
    
    # Find all YAML files (excluding backups)
    yaml_files = sorted(templates_dir.rglob("*.yaml"))
    yaml_files = [f for f in yaml_files if "backup" not in f.name]
    
    if not yaml_files:
        print("No YAML files found")
        sys.exit(1)
    
    updated_count = 0
    skipped_count = 0
    failed_count = 0
    
    for file_path in yaml_files:
        result = enhance_file(file_path)
        if result is True:
            updated_count += 1
        elif result is None:
            skipped_count += 1
        else:
            failed_count += 1
    
    print()
    print("==========================================")
    print("Summary")
    print("==========================================")
    print(f"✅ Updated: {updated_count} files")
    print(f"⏭️  Skipped: {skipped_count} files")
    if failed_count > 0:
        print(f"❌ Failed: {failed_count} files")
    print()
    print("Done!")


if __name__ == "__main__":
    main()