automation.py

"""
Automation Script for HuggingFace Space Management

This script handles automated synchronization between GitHub and HuggingFace Spaces.
Can be run as a standalone script or scheduled via cron/GitHub Actions.
"""

import os
import sys
import argparse
import yaml
from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv
from hf_space_sync import HFSpaceSync
from genesis_boiler import GenesisBoiler
import json

load_dotenv()


class SpaceAutomation:
    """Automate HuggingFace Space management tasks."""

    def __init__(self, config_path: str = "config.yaml", verbose: bool = True):
        """
        Initialize automation system.

        Args:
            config_path: Path to configuration file
            verbose: Enable verbose output
        """
        self.config_path = config_path
        self.verbose = verbose
        self.config = self._load_config()
        self.log_file = f"automation_log_{datetime.now().strftime('%Y%m%d')}.txt"

    def _load_config(self):
        """Load configuration."""
        try:
            with open(self.config_path, 'r') as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            self.log(f"ERROR: Config file {self.config_path} not found")
            sys.exit(1)

    def log(self, message: str):
        """Log a message."""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        log_msg = f"[{timestamp}] {message}"

        if self.verbose:
            print(log_msg)

        with open(self.log_file, 'a') as f:
            f.write(log_msg + "\n")

    def run_audit(self) -> dict:
        """
        Run file system audit.

        Returns:
            Audit results
        """
        self.log("Starting file audit...")

        try:
            boiler = GenesisBoiler(self.config_path)
            results = boiler.run_full_audit()

            self.log(f"Audit complete: {results['file_count']} files processed")
            self.log(f"Inventory: {results['inventory']}")
            self.log(f"Archive: {results['archive']}")

            return results
        except Exception as e:
            self.log(f"ERROR: Audit failed - {e}")
            raise

    def sync_all_spaces(self) -> dict:
        """
        Synchronize all configured spaces.

        Returns:
            Dictionary of sync results per space
        """
        self.log("Starting space synchronization...")

        try:
            hf_sync = HFSpaceSync(self.config_path)
            results = {}

            spaces = self.config.get('spaces', {})
            for space_key, space_config in spaces.items():
                if not space_config.get('auto_sync', False):
                    self.log(f"Skipping {space_key} (auto_sync disabled)")
                    continue

                space_name = space_config['name']
                self.log(f"Syncing {space_key} ({space_name})...")

                try:
                    sync_result = hf_sync.sync_directory(space_name, ".")
                    results[space_key] = sync_result

                    self.log(f"✓ {space_key}: {sync_result['uploaded']} uploaded, "
                           f"{sync_result['skipped']} skipped")
                except Exception as e:
                    self.log(f"ERROR: Failed to sync {space_key} - {e}")
                    results[space_key] = {"error": str(e)}

            return results
        except Exception as e:
            self.log(f"ERROR: Space sync failed - {e}")
            raise

    def sync_specific_space(self, space_name: str, local_dir: str = ".") -> dict:
        """
        Synchronize a specific space.

        Args:
            space_name: Name of the space to sync
            local_dir: Local directory to sync

        Returns:
            Sync results
        """
        self.log(f"Syncing {space_name} from {local_dir}...")

        try:
            hf_sync = HFSpaceSync(self.config_path)
            result = hf_sync.sync_directory(space_name, local_dir)

            self.log(f"✓ Sync complete: {result['uploaded']} uploaded, "
                   f"{result['skipped']} skipped")

            return result
        except Exception as e:
            self.log(f"ERROR: Sync failed - {e}")
            raise

    def backup_to_mapping_inventory(self) -> dict:
        """
        Backup current repository to Mapping-and-Inventory space.

        Returns:
            Backup results
        """
        self.log("Creating backup to Mapping-and-Inventory...")

        try:
            # First run audit to get current state
            audit_results = self.run_audit()

            # Get mapping-inventory space config
            mapping_space = self.config.get('spaces', {}).get('mapping_inventory', {})
            space_name = mapping_space.get('name', 'mapping-and-inventory')

            # Sync to mapping-inventory space
            hf_sync = HFSpaceSync(self.config_path)
            sync_result = hf_sync.sync_directory(space_name, ".")

            # Also upload the inventory and archive
            if audit_results.get('inventory'):
                hf_sync.upload_files(
                    space_name,
                    audit_results['inventory'],
                    commit_message="Automated inventory backup"
                )

            if audit_results.get('archive'):
                hf_sync.upload_files(
                    space_name,
                    audit_results['archive'],
                    commit_message="Automated archive backup"
                )

            self.log(f"✓ Backup complete to {space_name}")

            return {
                "audit": audit_results,
                "sync": sync_result
            }
        except Exception as e:
            self.log(f"ERROR: Backup failed - {e}")
            raise

    def create_summary_report(self, results: dict) -> str:
        """
        Create a summary report of automation results.

        Args:
            results: Dictionary of results from automation tasks

        Returns:
            Path to summary report file
        """
        report_path = f"automation_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"

        report = {
            "timestamp": datetime.now().isoformat(),
            "config_file": self.config_path,
            "results": results,
            "log_file": self.log_file
        }

        with open(report_path, 'w') as f:
            json.dump(report, f, indent=2)

        self.log(f"Summary report created: {report_path}")
        return report_path


def main():
    """Main entry point for automation script."""
    parser = argparse.ArgumentParser(
        description="VAMGUARD TITAN - HuggingFace Space Automation"
    )

    parser.add_argument(
        'command',
        choices=['audit', 'sync', 'sync-space', 'backup', 'full'],
        help='Command to execute'
    )

    parser.add_argument(
        '--space',
        help='Space name (for sync-space command)'
    )

    parser.add_argument(
        '--dir',
        default='.',
        help='Local directory to sync (default: current directory)'
    )

    parser.add_argument(
        '--config',
        default='config.yaml',
        help='Path to config file (default: config.yaml)'
    )

    parser.add_argument(
        '--quiet',
        action='store_true',
        help='Suppress verbose output'
    )

    args = parser.parse_args()

    # Initialize automation
    automation = SpaceAutomation(args.config, verbose=not args.quiet)

    try:
        results = {}

        if args.command == 'audit':
            results['audit'] = automation.run_audit()

        elif args.command == 'sync':
            results['sync'] = automation.sync_all_spaces()

        elif args.command == 'sync-space':
            if not args.space:
                automation.log("ERROR: --space required for sync-space command")
                sys.exit(1)
            results['sync_space'] = automation.sync_specific_space(args.space, args.dir)

        elif args.command == 'backup':
            results['backup'] = automation.backup_to_mapping_inventory()

        elif args.command == 'full':
            automation.log("Running full automation workflow...")
            results['audit'] = automation.run_audit()
            results['sync'] = automation.sync_all_spaces()
            results['backup'] = automation.backup_to_mapping_inventory()

        # Create summary report
        report_path = automation.create_summary_report(results)
        automation.log(f"Automation complete! Report: {report_path}")

    except Exception as e:
        automation.log(f"FATAL ERROR: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()