#!/usr/bin/env python3
"""
Home Assistant Legacy Template Migration Script
Automatically converts legacy platform: template syntax to modern template: syntax
"""

import yaml
import os
import re
import shutil
from datetime import datetime
from pathlib import Path

# Configuration
CONFIG_DIR = Path("/config")
BACKUP_TIMESTAMP = datetime.now().strftime("%Y%m%d_%H%M%S")
BACKUP_DIR = CONFIG_DIR / "backups" / f"template_migration_{BACKUP_TIMESTAMP}"

# Colors for output
GREEN = '\033[0;32m'
BLUE = '\033[0;34m'
YELLOW = '\033[1;33m'
RED = '\033[0;31m'
NC = '\033[0m'

class TemplateMigrator:
    def __init__(self):
        self.legacy_sensors = []
        self.legacy_binary_sensors = []
        self.files_to_update = {}
        
    def find_legacy_templates(self):
        """Find all legacy template definitions in YAML files"""
        print(f"{BLUE}ℹ{NC} Searching for legacy template definitions...")
        
        # Search common locations
        yaml_files = []
        for pattern in ['*.yaml', '*.yml']:
            yaml_files.extend(CONFIG_DIR.glob(pattern))
            yaml_files.extend((CONFIG_DIR / 'packages').glob(f'**/{pattern}'))
        
        for yaml_file in yaml_files:
            try:
                with open(yaml_file, 'r') as f:
                    content = f.read()
                    
                # Check for legacy sensor templates
                if re.search(r'^\s*-?\s*platform:\s*template', content, re.MULTILINE):
                    self.analyze_file(yaml_file)
            except Exception as e:
                print(f"{YELLOW}⚠{NC} Error reading {yaml_file}: {e}")
    
    def analyze_file(self, filepath):
        """Analyze a YAML file for legacy templates"""
        try:
            with open(filepath, 'r') as f:
                data = yaml.safe_load(f)
            
            if not data:
                return
            
            file_info = {
                'path': filepath,
                'sensors': [],
                'binary_sensors': []
            }
            
            # Check for legacy sensor templates
            if 'sensor' in data and isinstance(data['sensor'], list):
                for sensor in data['sensor']:
                    if isinstance(sensor, dict) and sensor.get('platform') == 'template':
                        if 'sensors' in sensor:
                            for name, config in sensor['sensors'].items():
                                file_info['sensors'].append({
                                    'name': name,
                                    'config': config
                                })
            
            # Check for legacy binary_sensor templates
            if 'binary_sensor' in data and isinstance(data['binary_sensor'], list):
                for sensor in data['binary_sensor']:
                    if isinstance(sensor, dict) and sensor.get('platform') == 'template':
                        if 'sensors' in sensor:
                            for name, config in sensor['sensors'].items():
                                file_info['binary_sensors'].append({
                                    'name': name,
                                    'config': config
                                })
            
            if file_info['sensors'] or file_info['binary_sensors']:
                self.files_to_update[str(filepath)] = file_info
                print(f"  {GREEN}✓{NC} Found legacy templates in: {filepath.name}")
                if file_info['sensors']:
                    print(f"    - {len(file_info['sensors'])} sensor(s)")
                if file_info['binary_sensors']:
                    print(f"    - {len(file_info['binary_sensors'])} binary_sensor(s)")
                
        except Exception as e:
            print(f"{YELLOW}⚠{NC} Error analyzing {filepath}: {e}")
    
    def convert_to_modern_syntax(self, name, config, sensor_type='sensor'):
        """Convert legacy template config to modern syntax"""
        modern = {}
        
        # Required fields
        modern['name'] = config.get('friendly_name', name.replace('_', ' ').title())
        modern['state'] = config.get('value_template', '')
        
        # Optional fields
        if 'icon_template' in config or 'icon' in config:
            modern['icon'] = config.get('icon_template', config.get('icon', ''))
        
        if 'entity_picture_template' in config or 'entity_picture' in config:
            modern['picture'] = config.get('entity_picture_template', config.get('entity_picture', ''))
        
        if 'availability_template' in config or 'availability' in config:
            modern['availability'] = config.get('availability_template', config.get('availability', ''))
        
        if 'attribute_templates' in config:
            modern['attributes'] = config['attribute_templates']
        
        # Sensor-specific fields
        if sensor_type == 'sensor':
            if 'unit_of_measurement' in config:
                modern['unit_of_measurement'] = config['unit_of_measurement']
            if 'device_class' in config:
                modern['device_class'] = config['device_class']
            if 'state_class' in config:
                modern['state_class'] = config['state_class']
        
        # Binary sensor specific fields
        if sensor_type == 'binary_sensor':
            if 'device_class' in config:
                modern['device_class'] = config['device_class']
        
        # Entity ID preservation
        modern['unique_id'] = name
        
        return modern
    
    def generate_modern_template_yaml(self):
        """Generate the modern template: section YAML"""
        template_config = []
        
        # Process all sensors
        all_sensors = []
        all_binary_sensors = []
        
        for file_path, info in self.files_to_update.items():
            for sensor in info['sensors']:
                all_sensors.append(
                    self.convert_to_modern_syntax(sensor['name'], sensor['config'], 'sensor')
                )
            
            for binary_sensor in info['binary_sensors']:
                all_binary_sensors.append(
                    self.convert_to_modern_syntax(binary_sensor['name'], binary_sensor['config'], 'binary_sensor')
                )
        
        # Build template structure
        if all_sensors:
            template_config.append({'sensor': all_sensors})
        
        if all_binary_sensors:
            template_config.append({'binary_sensor': all_binary_sensors})
        
        return template_config
    
    def backup_files(self):
        """Backup all files that will be modified"""
        print(f"\n{BLUE}ℹ{NC} Creating backups...")
        BACKUP_DIR.mkdir(parents=True, exist_ok=True)
        
        for file_path in self.files_to_update.keys():
            src = Path(file_path)
            dst = BACKUP_DIR / src.name
            shutil.copy2(src, dst)
            print(f"  {GREEN}✓{NC} Backed up: {src.name}")
    
    def remove_legacy_definitions(self):
        """Remove legacy template definitions from files"""
        print(f"\n{BLUE}ℹ{NC} Removing legacy template definitions...")
        
        for file_path, info in self.files_to_update.items():
            try:
                with open(file_path, 'r') as f:
                    data = yaml.safe_load(f)
                
                modified = False
                
                # Remove legacy sensor templates
                if 'sensor' in data and isinstance(data['sensor'], list):
                    data['sensor'] = [
                        s for s in data['sensor']
                        if not (isinstance(s, dict) and s.get('platform') == 'template')
                    ]
                    if not data['sensor']:
                        del data['sensor']
                    modified = True
                
                # Remove legacy binary_sensor templates
                if 'binary_sensor' in data and isinstance(data['binary_sensor'], list):
                    data['binary_sensor'] = [
                        s for s in data['binary_sensor']
                        if not (isinstance(s, dict) and s.get('platform') == 'template')
                    ]
                    if not data['binary_sensor']:
                        del data['binary_sensor']
                    modified = True
                
                if modified:
                    with open(file_path, 'w') as f:
                        yaml.dump(data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
                    print(f"  {GREEN}✓{NC} Updated: {Path(file_path).name}")
                    
            except Exception as e:
                print(f"{RED}✗{NC} Error updating {file_path}: {e}")
    
    def add_modern_templates(self, template_yaml):
        """Add modern template definitions to configuration.yaml"""
        print(f"\n{BLUE}ℹ{NC} Adding modern template definitions...")
        
        config_file = CONFIG_DIR / "configuration.yaml"
        
        try:
            with open(config_file, 'r') as f:
                config = yaml.safe_load(f) or {}
            
            # Check if template: already exists
            if 'template' in config:
                print(f"  {YELLOW}ℹ{NC} Existing template: section found, merging...")
                existing_template = config['template']
                if not isinstance(existing_template, list):
                    existing_template = [existing_template]
                
                # Merge new templates
                for new_template in template_yaml:
                    existing_template.append(new_template)
                
                config['template'] = existing_template
            else:
                config['template'] = template_yaml
            
            # Write updated configuration
            with open(config_file, 'w') as f:
                yaml.dump(config, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
            
            print(f"  {GREEN}✓{NC} Added modern template definitions to configuration.yaml")
            
        except Exception as e:
            print(f"{RED}✗{NC} Error updating configuration.yaml: {e}")
            raise
    
    def create_restore_script(self):
        """Create a restore script in the backup directory"""
        restore_script = BACKUP_DIR / "restore.sh"
        
        with open(restore_script, 'w') as f:
            f.write(f"""#!/bin/bash
echo "Restoring template migration backup..."
cp "{BACKUP_DIR}"/*.yaml "{CONFIG_DIR}/"
echo "Files restored!"
echo "Restart Home Assistant: ha core restart"
""")
        
        restore_script.chmod(0o755)
        print(f"\n{GREEN}✓{NC} Restore script created: {restore_script}")

def main():
    print("=" * 70)
    print("Home Assistant Legacy Template Migration")
    print("=" * 70)
    print()
    
    migrator = TemplateMigrator()
    
    # Step 1: Find all legacy templates
    migrator.find_legacy_templates()
    
    if not migrator.files_to_update:
        print(f"\n{GREEN}✓{NC} No legacy template definitions found!")
        return
    
    # Count total entities
    total_sensors = sum(len(info['sensors']) for info in migrator.files_to_update.values())
    total_binary = sum(len(info['binary_sensors']) for info in migrator.files_to_update.values())
    total = total_sensors + total_binary
    
    print(f"\n{BLUE}ℹ{NC} Found {total} legacy template entities:")
    print(f"  - {total_sensors} sensor(s)")
    print(f"  - {total_binary} binary_sensor(s)")
    
    # Step 2: Generate modern template YAML
    modern_template = migrator.generate_modern_template_yaml()
    
    # Step 3: Create backups
    migrator.backup_files()
    
    # Step 4: Remove legacy definitions
    migrator.remove_legacy_definitions()
    
    # Step 5: Add modern template definitions
    migrator.add_modern_templates(modern_template)
    
    # Step 6: Create restore script
    migrator.create_restore_script()
    
    print("\n" + "=" * 70)
    print("MIGRATION COMPLETE")
    print("=" * 70)
    print(f"\nBackup location: {BACKUP_DIR}")
    print(f"Migrated: {total} template entities")
    print()
    print("NEXT STEPS:")
    print("1. Review configuration.yaml for the new template: section")
    print("2. Check config: ha core check")
    print("3. Restart HA: ha core restart")
    print(f"4. If issues: bash {BACKUP_DIR}/restore.sh")
    print()
    print("=" * 70)

if __name__ == "__main__":
    main()