ha-chore-cleanup/ha_chore_cleanup.py

#!/usr/bin/env python3
"""
Home Assistant Chore Cleanup Script
Removes unused chore-related entities from YAML configuration files
Creates backups before making any changes
"""

import os
import yaml
import shutil
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Set, Tuple

# Entities to remove
INPUT_BOOLEANS_TO_REMOVE = {
    'task_lou_clean_desk_pending',
    'task_jess_clean_desk_pending',
    'task_william_clean_desk_pending',
    'task_xander_clean_desk_pending',
    'task_bella_clean_desk_pending',
    'task_lou_clean_room_pending',
    'task_jess_clean_room_pending',
    'task_william_clean_room_pending',
    'task_xander_clean_room_pending',
    'task_bella_clean_room_pending',
    'task_tidy_lounge_pending',
    'task_vacuum_room_pending',
    'task_get_school_ready_pending',
    'task_clean_desks_done_this_week',
    'task_clean_rooms_done_this_week',
    'task_tidy_lounge_done_this_week',
    'task_tidy_kitchen_done_today',
    'task_clean_dining_table_done_today',
    'task_vacuum_room_done_today',
}

SCRIPTS_TO_REMOVE = {
    'complete_task_dishwasher_unload',
    'complete_task_washing_machine_unload',
    'complete_task_dryer_unload',
    'complete_task_bins_out',
    'complete_task_bins_in',
    'complete_task_kitty_litter_clean',
    'complete_task_kitty_litter_change',
    'complete_tidy_lounge',
    'complete_vacuum_room',
    'complete_get_school_ready',
}


class YAMLCleaner:
    def __init__(self, config_dir: str = "/config"):
        self.config_dir = Path(config_dir)
        self.backup_dir = self.config_dir / "backups" / f"cleanup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.removed_items: Dict[str, List[str]] = {
            'input_booleans': [],
            'scripts': []
        }
        self.errors: List[str] = []
        
    def create_backup_dir(self):
        """Create backup directory"""
        self.backup_dir.mkdir(parents=True, exist_ok=True)
        print(f"✓ Created backup directory: {self.backup_dir}")
        
    def backup_file(self, file_path: Path) -> Path:
        """Create backup of a file before modifying"""
        if not file_path.exists():
            return None
            
        backup_path = self.backup_dir / file_path.name
        shutil.copy2(file_path, backup_path)
        print(f"  ✓ Backed up: {file_path.name}")
        return backup_path
        
    def load_yaml_file(self, file_path: Path) -> Tuple[Dict, bool]:
        """Load YAML file with proper handling"""
        if not file_path.exists():
            return None, False
            
        try:
            with open(file_path, 'r') as f:
                content = f.read()
                
            # Check if file uses !include or other special directives
            has_directives = '!include' in content or '!secret' in content
            
            if has_directives:
                # For files with directives, we'll need to be more careful
                return content, True
            else:
                data = yaml.safe_load(content)
                return data, False
                
        except Exception as e:
            self.errors.append(f"Error loading {file_path.name}: {str(e)}")
            return None, False
            
    def save_yaml_file(self, file_path: Path, data, is_raw: bool = False):
        """Save YAML file"""
        try:
            with open(file_path, 'w') as f:
                if is_raw:
                    f.write(data)
                else:
                    yaml.dump(data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
            print(f"  ✓ Saved: {file_path.name}")
            return True
        except Exception as e:
            self.errors.append(f"Error saving {file_path.name}: {str(e)}")
            return False
            
    def clean_input_booleans_from_config(self, config_data: Dict) -> Dict:
        """Remove input_boolean entries from configuration.yaml structure"""
        if not config_data or not isinstance(config_data, dict):
            return config_data
            
        if 'input_boolean' in config_data:
            original_count = len(config_data['input_boolean']) if isinstance(config_data['input_boolean'], dict) else 0
            
            if isinstance(config_data['input_boolean'], dict):
                for key in list(config_data['input_boolean'].keys()):
                    if key in INPUT_BOOLEANS_TO_REMOVE:
                        del config_data['input_boolean'][key]
                        self.removed_items['input_booleans'].append(key)
                        
            new_count = len(config_data['input_boolean']) if isinstance(config_data['input_boolean'], dict) else 0
            print(f"  Removed {original_count - new_count} input_boolean entries")
            
        return config_data
        
    def clean_scripts_from_config(self, config_data: Dict) -> Dict:
        """Remove script entries from configuration.yaml structure"""
        if not config_data or not isinstance(config_data, dict):
            return config_data
            
        if 'script' in config_data:
            original_count = len(config_data['script']) if isinstance(config_data['script'], dict) else 0
            
            if isinstance(config_data['script'], dict):
                for key in list(config_data['script'].keys()):
                    if key in SCRIPTS_TO_REMOVE:
                        del config_data['script'][key]
                        self.removed_items['scripts'].append(key)
                        
            new_count = len(config_data['script']) if isinstance(config_data['script'], dict) else 0
            print(f"  Removed {original_count - new_count} script entries")
            
        return config_data
        
    def clean_standalone_file(self, file_path: Path, entity_type: str):
        """Clean a standalone YAML file (input_boolean.yaml or scripts.yaml)"""
        print(f"\nProcessing {file_path.name}...")
        
        data, is_raw = self.load_yaml_file(file_path)
        if data is None:
            print(f"  ! File not found or empty, skipping")
            return
            
        # Backup original
        self.backup_file(file_path)
        
        if is_raw:
            print(f"  ! File contains directives (!include, etc.), manual review required")
            return
            
        if not isinstance(data, dict):
            print(f"  ! Unexpected format, skipping")
            return
            
        # Remove items
        original_count = len(data)
        to_remove = INPUT_BOOLEANS_TO_REMOVE if entity_type == 'input_boolean' else SCRIPTS_TO_REMOVE
        
        for key in list(data.keys()):
            if key in to_remove:
                del data[key]
                self.removed_items[entity_type + 's'].append(key)
                
        new_count = len(data)
        removed_count = original_count - new_count
        
        if removed_count > 0:
            self.save_yaml_file(file_path, data)
            print(f"  ✓ Removed {removed_count} {entity_type} entries")
        else:
            print(f"  No items to remove")
            
    def process_configuration_yaml(self):
        """Process main configuration.yaml file"""
        config_path = self.config_dir / "configuration.yaml"
        print(f"\nProcessing configuration.yaml...")
        
        data, is_raw = self.load_yaml_file(config_path)
        if data is None:
            print(f"  ! File not found or empty")
            return
            
        if is_raw:
            print(f"  ! File contains directives, will check for standalone files instead")
            return
            
        # Backup original
        self.backup_file(config_path)
        
        # Clean both types
        data = self.clean_input_booleans_from_config(data)
        data = self.clean_scripts_from_config(data)
        
        # Save
        self.save_yaml_file(config_path, data)
        
    def find_config_files(self) -> Dict[str, List[Path]]:
        """Find all potential configuration files"""
        files = {
            'input_boolean': [],
            'script': []
        }
        
        # Check for standalone files
        for pattern in ['input_boolean.yaml', 'input_booleans.yaml', 'helpers.yaml']:
            matches = list(self.config_dir.glob(pattern))
            files['input_boolean'].extend(matches)
            
        for pattern in ['scripts.yaml', 'script.yaml']:
            matches = list(self.config_dir.glob(pattern))
            files['script'].extend(matches)
            
        # Check packages directory
        packages_dir = self.config_dir / "packages"
        if packages_dir.exists():
            for yaml_file in packages_dir.glob("**/*.yaml"):
                # Check if file might contain our entities
                try:
                    with open(yaml_file, 'r') as f:
                        content = f.read().lower()
                        if 'input_boolean' in content and any(key in content for key in INPUT_BOOLEANS_TO_REMOVE):
                            files['input_boolean'].append(yaml_file)
                        if 'script' in content and any(key in content for key in SCRIPTS_TO_REMOVE):
                            files['script'].append(yaml_file)
                except:
                    pass
                    
        return files
        
    def run_cleanup(self):
        """Execute the cleanup process"""
        print("=" * 70)
        print("Home Assistant Chore Cleanup Script")
        print("=" * 70)
        
        # Create backup directory
        self.create_backup_dir()
        
        # Process main configuration file
        self.process_configuration_yaml()
        
        # Find and process standalone files
        config_files = self.find_config_files()
        
        # Process input_boolean files
        for file_path in config_files['input_boolean']:
            self.clean_standalone_file(file_path, 'input_boolean')
            
        # Process script files
        for file_path in config_files['script']:
            self.clean_standalone_file(file_path, 'script')
            
        # Print summary
        self.print_summary()
        
    def print_summary(self):
        """Print cleanup summary"""
        print("\n" + "=" * 70)
        print("CLEANUP SUMMARY")
        print("=" * 70)
        
        print(f"\n✓ Backups saved to: {self.backup_dir}")
        
        print(f"\n📋 Input Booleans Removed: {len(self.removed_items['input_booleans'])}")
        for item in sorted(self.removed_items['input_booleans']):
            print(f"  - {item}")
            
        print(f"\n📜 Scripts Removed: {len(self.removed_items['scripts'])}")
        for item in sorted(self.removed_items['scripts']):
            print(f"  - {item}")
            
        if self.errors:
            print(f"\n⚠️  Errors encountered: {len(self.errors)}")
            for error in self.errors:
                print(f"  - {error}")
                
        print("\n" + "=" * 70)
        print("NEXT STEPS:")
        print("=" * 70)
        print("1. Review the changes in your YAML files")
        print("2. Run 'ha core check' to verify configuration")
        print("3. If everything looks good, restart Home Assistant")
        print("4. If issues occur, restore from backups in:")
        print(f"   {self.backup_dir}")
        print("=" * 70)
        
        # Create a restoration script
        self.create_restore_script()
        
    def create_restore_script(self):
        """Create a script to restore from backups if needed"""
        restore_script = self.backup_dir / "restore.sh"
        
        script_content = f"""#!/bin/bash
# Restore script - Run this if you need to undo the cleanup

echo "Restoring files from backup..."

"""
        
        for backup_file in self.backup_dir.glob("*.yaml"):
            original = self.config_dir / backup_file.name
            script_content += f'cp "{backup_file}" "{original}"\n'
            
        script_content += """
echo "Files restored!"
echo "Please restart Home Assistant for changes to take effect"
"""
        
        with open(restore_script, 'w') as f:
            f.write(script_content)
            
        restore_script.chmod(0o755)
        print(f"\n✓ Created restore script: {restore_script}")


def main():
    """Main entry point"""
    import sys
    
    # Check if running in Home Assistant environment
    config_dir = "/config"
    
    # Allow override via command line
    if len(sys.argv) > 1:
        config_dir = sys.argv[1]
        
    if not os.path.exists(config_dir):
        print(f"Error: Config directory not found: {config_dir}")
        print("Usage: python3 ha_chore_cleanup.py [config_directory]")
        print("Default: /config")
        sys.exit(1)
        
    cleaner = YAMLCleaner(config_dir)
    cleaner.run_cleanup()


if __name__ == "__main__":
    main()