src/data_sync.py

#!/usr/bin/env python3
"""
Hermes Agent Data Sync Service
Handles data persistence to/from Hugging Face Dataset
"""
import os
import sys
import time
import json
import shutil
import tarfile
import argparse
import hashlib
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, List

from huggingface_hub import HfApi, hf_hub_download, upload_folder
from loguru import logger
import re

# File watching (optional)
try:
    from watchdog.observers import Observer
    from watchdog.events import FileSystemEventHandler
    WATCHDOG_AVAILABLE = True
except ImportError:
    WATCHDOG_AVAILABLE = False
    logger.warning("watchdog not installed, file change detection disabled")


class DatasetManager:
    """Manages data synchronization with Hugging Face Dataset"""
    
    def __init__(self, dataset_repo: Optional[str] = None, token: Optional[str] = None):
        self.dataset_repo = dataset_repo or os.environ.get('HF_DATASET_REPO')
        self.token = token or os.environ.get('HF_TOKEN') or os.environ.get('HUGGING_FACE_HUB_TOKEN')
        self.api = HfApi(token=self.token)
        self.hermes_home = Path(os.environ.get('HERMES_HOME', '/data/.hermes'))
        self.temp_dir = Path('/tmp/hermes_sync')
        self.state_file = self.hermes_home / '.sync_state.json'
        
        # Incremental backup config (controllable via environment variables)
        self.enable_incremental = os.environ.get('SYNC_INCREMENTAL', 'true').lower() in ('true', '1', 'yes')
        self.sync_interval_override = os.environ.get('SYNC_INTERVAL', None)
        
        # Data path mapping
        # skip_backup=True means skip this path in incremental mode (large files / temp files)
        self.path_mapping = {
            'config': {'path': self.hermes_home / 'config.yaml', 'skip_backup': False},
            'env': {'path': self.hermes_home / '.env', 'skip_backup': False},
            'auth': {'path': self.hermes_home / 'auth.json', 'skip_backup': False},
            'soul': {'path': self.hermes_home / 'SOUL.md', 'skip_backup': False},
            'memories': {'path': self.hermes_home / 'memories', 'skip_backup': False},
            'skills': {'path': self.hermes_home / 'skills', 'skip_backup': False},
            'sessions': {'path': self.hermes_home / 'sessions', 'skip_backup': False},
            'state_db': {'path': self.hermes_home / 'state.db', 'skip_backup': False},
            'logs': {'path': self.hermes_home / 'logs', 'skip_backup': True},  # logs are large and not critical
            'cron': {'path': self.hermes_home / 'cron', 'skip_backup': False},
            'webui_token': {'path': Path('/data/.hermes-web-ui') / '.token', 'skip_backup': False},
            'webui_db': {'path': Path('/home/appuser/.hermes-web-ui'), 'skip_backup': False},  # Web UI account DB + sessions
            'image_cache': {'path': self.hermes_home / 'image_cache', 'skip_backup': True},  # image cache is large
            'baoyu_skills': {'path': Path('/home/appuser/.baoyu-skills'), 'skip_backup': False},
        }
        
    def validate(self) -> bool:
        """Validate configuration"""
        if not self.dataset_repo:
            logger.error("HF_DATASET_REPO not set")
            return False
        
        if not self.token:
            logger.warning("HF_TOKEN not set, will try public dataset")
            
        return True
    
    def _calculate_file_hash(self, file_path: Path) -> str:
        """Calculate file SHA256 hash (reads first 128KB only for speed)"""
        if not file_path.exists():
            return ""
        try:
            sha256 = hashlib.sha256()
            with open(file_path, 'rb') as f:
                data = f.read(131072)  # 128KB
                sha256.update(data)
            return sha256.hexdigest()
        except Exception:
            return ""
    
    def _calculate_dir_hash(self, dir_path: Path) -> str:
        """Calculate combined hash of all files in a directory"""
        if not dir_path.exists():
            return ""
        try:
            sha256 = hashlib.sha256()
            for root, dirs, files in os.walk(dir_path):
                # Skip hidden files
                files = [f for f in files if not f.startswith('.')]
                for f in sorted(files):
                    file_path = Path(root) / f
                    if file_path.exists():
                        try:
                            with open(file_path, 'rb') as fp:
                                data = fp.read(65536)  # 64KB
                                sha256.update(data)
                        except Exception:
                            pass
            return sha256.hexdigest()
        except Exception:
            return ""
    
    def _load_sync_state(self) -> Dict:
        """Load last sync state"""
        if self.state_file.exists():
            try:
                with open(self.state_file, 'r') as f:
                    return json.load(f)
            except Exception:
                pass
        return {'files': {}, 'dirs': {}, 'last_sync': 0, 'version': 1}
    
    def _save_sync_state(self, state: Dict):
        """Save sync state"""
        try:
            self.state_file.parent.mkdir(parents=True, exist_ok=True)
            with open(self.state_file, 'w') as f:
                json.dump(state, f, indent=2)
        except Exception as e:
            logger.warning(f"Failed to save sync state: {e}")
    
    def _has_changes(self, state: Dict) -> bool:
        """Check for file changes (core of incremental backup logic)"""
        if not state.get('last_sync'):
            return True  # First sync
        
        for key, info in self.path_mapping.items():
            if info.get('skip_backup'):
                continue
            path = info['path']
            if path.is_file():
                current_hash = self._calculate_file_hash(path)
                old_hash = state.get('files', {}).get(key, '')
                if current_hash and current_hash != old_hash:
                    logger.debug(f"File changed: {key}")
                    return True
            elif path.is_dir():
                current_hash = self._calculate_dir_hash(path)
                old_hash = state.get('dirs', {}).get(key, '')
                if current_hash and current_hash != old_hash:
                    logger.debug(f"Dir changed: {key}")
                    return True
        
        return False
    
    def prepare_backup_data(self) -> Path:
        """Prepare backup data into temp directory (incremental backup supported)"""
        logger.info("Preparing backup data...")
        
        # Incremental backup check - skip if no changes
        if self.enable_incremental:
            state = self._load_sync_state()
            if not self._has_changes(state):
                logger.info("No changes detected, skipping backup (incremental mode)")
                # Update last sync time
                state['last_sync'] = time.time()
                self._save_sync_state(state)
                return Path('/dev/null')  # Sentinel value  -  upload_to_dataset will detect this and skip
        
        # Clean up and create temp directory
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
        self.temp_dir.mkdir(parents=True)
        
        # Create directory structure (logs/ and image_cache/ intentionally excluded)
        (self.temp_dir / 'config').mkdir()
        (self.temp_dir / 'personality').mkdir()
        (self.temp_dir / 'memories').mkdir()
        (self.temp_dir / 'skills').mkdir()
        (self.temp_dir / 'sessions').mkdir()
        (self.temp_dir / 'state').mkdir()
        (self.temp_dir / 'cron').mkdir()
        (self.temp_dir / 'webui').mkdir()
        (self.temp_dir / 'baoyu_skills').mkdir()
        
        # Copy files
        try:
            # Config file
            config_info = self.path_mapping['config']
            if not config_info.get('skip_backup') and config_info['path'].exists():
                shutil.copy2(config_info['path'], self.temp_dir / 'config' / 'config.yaml')
            
            # Environment variables (sensitive) - mask before uploading
            env_info = self.path_mapping['env']
            if not env_info.get('skip_backup') and env_info['path'].exists():
                env_src = env_info['path']
                env_dst = self.temp_dir / 'config' / '.env'
                with open(env_src, 'r', encoding='utf-8') as f:
                    env_content = f.read()
                # Mask HF tokens and ALL other API keys before uploading
                env_content = re.sub(r'(HF_TOKEN|HUGGING_FACE_HUB_TOKEN)=hf_[a-zA-Z0-9]+', r'\1=***masked***', env_content)
                env_content = re.sub(r'(HF_TOKEN|HUGGING_FACE_HUB_TOKEN)=.+', r'\1=***masked***', env_content)
                env_content = re.sub(r'([A-Z_]*(API_KEY|SECRET|TOKEN)[A-Z_]*)=\S+', r'\1=***masked***', env_content)
                with open(env_dst, 'w', encoding='utf-8') as f:
                    f.write(env_content)
            
            # OAuth credentials
            auth_info = self.path_mapping['auth']
            if not auth_info.get('skip_backup') and auth_info['path'].exists():
                shutil.copy2(auth_info['path'], self.temp_dir / 'config' / 'auth.json')
            
            # Agent personality
            soul_info = self.path_mapping['soul']
            if not soul_info.get('skip_backup') and soul_info['path'].exists():
                shutil.copy2(soul_info['path'], self.temp_dir / 'personality' / 'SOUL.md')
            
            # Memories
            mem_info = self.path_mapping['memories']
            if not mem_info.get('skip_backup') and mem_info['path'].exists():
                shutil.copytree(mem_info['path'], self.temp_dir / 'memories', dirs_exist_ok=True)
            
            # Skills
            skill_info = self.path_mapping['skills']
            if not skill_info.get('skip_backup') and skill_info['path'].exists():
                shutil.copytree(skill_info['path'], self.temp_dir / 'skills', dirs_exist_ok=True)
            
            # Sessions
            sess_info = self.path_mapping['sessions']
            if not sess_info.get('skip_backup') and sess_info['path'].exists():
                shutil.copytree(sess_info['path'], self.temp_dir / 'sessions', dirs_exist_ok=True)
            
            # Database
            db_info = self.path_mapping['state_db']
            if not db_info.get('skip_backup') and db_info['path'].exists():
                shutil.copy2(db_info['path'], self.temp_dir / 'state' / 'state.db')
            
            # Logs - skipped (large and not critical)
            # if self.path_mapping['logs'].exists():
            #     shutil.copytree(self.path_mapping['logs'], self.temp_dir / 'logs', dirs_exist_ok=True)
            
            # Scheduled tasks (cron)
            cron_info = self.path_mapping['cron']
            if not cron_info.get('skip_backup') and cron_info['path'].exists():
                shutil.copytree(cron_info['path'], self.temp_dir / 'cron', dirs_exist_ok=True)
            
            # Image cache - skipped (large files)
            # if self.path_mapping['image_cache'].exists():
            #     shutil.copytree(self.path_mapping['image_cache'], self.temp_dir / 'image_cache', dirs_exist_ok=True)
            
            # baoyu-skills user config
            baoyu_info = self.path_mapping['baoyu_skills']
            if not baoyu_info.get('skip_backup') and baoyu_info['path'].exists():
                shutil.copytree(baoyu_info['path'], self.temp_dir / 'baoyu_skills', dirs_exist_ok=True)
            
            # WebUI auth token
            webui_info = self.path_mapping['webui_token']
            if not webui_info.get('skip_backup') and webui_info['path'].exists():
                (self.temp_dir / 'webui').mkdir(exist_ok=True)
                shutil.copy2(webui_info['path'], self.temp_dir / 'webui' / '.token')

            # WebUI database (accounts, passwords, sessions, settings)
            webui_db_info = self.path_mapping['webui_db']
            if not webui_db_info.get('skip_backup') and webui_db_info['path'].exists():
                webui_db_dst = self.temp_dir / 'webui' / 'db'
                shutil.copytree(webui_db_info['path'], webui_db_dst, dirs_exist_ok=True,
                                ignore=shutil.ignore_patterns('logs', '*.log'))

            # Add metadata
            metadata = {
                'timestamp': datetime.now().isoformat(),
                'version': '0.10.0',
                'hermes_home': str(self.hermes_home),
                'incremental': self.enable_incremental,
            }
            with open(self.temp_dir / 'metadata.json', 'w') as f:
                json.dump(metadata, f, indent=2)
            
            logger.success(f"Backup prepared at {self.temp_dir}")
            return self.temp_dir
            
        except Exception as e:
            logger.error(f"Failed to prepare backup: {e}")
            raise
    
    def upload_to_dataset(self, force: bool = False) -> bool:
        """Upload data to Hugging Face Dataset"""
        try:
            backup_dir = self.prepare_backup_data()
            
            # Check if this is the incremental backup skip sentinel
            if backup_dir == Path('/dev/null'):
                logger.info("Backup skipped (no changes)")
                return True
            
            logger.info(f"Uploading to dataset: {self.dataset_repo}")
            
            # Upload folder to dataset
            self.api.upload_folder(
                folder_path=str(backup_dir),
                repo_id=self.dataset_repo,
                repo_type="dataset",
                commit_message=f"Hermes Agent backup - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
            )
            
            # Update sync state
            if self.enable_incremental:
                state = self._load_sync_state()
                state['last_sync'] = time.time()
                # Update file hashes
                for key, info in self.path_mapping.items():
                    if info.get('skip_backup'):
                        continue
                    path = info['path']
                    if path.is_file():
                        state.setdefault('files', {})[key] = self._calculate_file_hash(path)
                    elif path.is_dir():
                        state.setdefault('dirs', {})[key] = self._calculate_dir_hash(path)
                self._save_sync_state(state)
            
            logger.success("Backup uploaded successfully")
            return True
            
        except Exception as e:
            logger.error(f"Failed to upload to dataset: {e}")
            return False
    
    def download_from_dataset(self) -> bool:
        """Download data from Hugging Face Dataset"""
        try:
            logger.info(f"Downloading from dataset: {self.dataset_repo}")
            
            # Create temp download directory
            download_dir = Path('/tmp/hermes_download')
            if download_dir.exists():
                shutil.rmtree(download_dir)
            download_dir.mkdir(parents=True)
            
            # Download all files
            self.api.snapshot_download(
                repo_id=self.dataset_repo,
                repo_type="dataset",
                local_dir=str(download_dir)
            )
            
            logger.success("Download completed")
            
            # Restore data to Hermes directory
            self.restore_from_download(download_dir)
            return True
            
        except Exception as e:
            logger.error(f"Failed to download from dataset: {e}")
            return False
    
    def restore_from_download(self, download_dir: Path):
        """Restore data from a downloaded directory.
        
        Note: config.yaml is skipped during restore because entrypoint.sh regenerates
        it from environment variables. Restoring an old config.yaml would overwrite
        the correct model config (e.g. replace current model with an old one).
        """
        logger.info("Restoring data to Hermes home...")
        
        # Ensure destination directory exists
        self.hermes_home.mkdir(parents=True, exist_ok=True)
        
        # Restore strategy control
        skip_restore = os.environ.get('SKIP_CONFIG_RESTORE', 'true').lower() in ('true', '1', 'yes')
        
        if skip_restore:
            skip_paths = {
                'config/config.yaml',
                'baoyu_skills',
            }
            logger.info(f"SKIP_CONFIG_RESTORE=true, skipping: {', '.join(skip_paths)}")
        else:
            skip_paths = set()
            logger.info("SKIP_CONFIG_RESTORE=false, restoring all backed-up configurations")
        
        restore_mapping = {
            'config/.env': self.path_mapping['env']['path'],
            'config/auth.json': self.path_mapping['auth']['path'],
            'personality/SOUL.md': self.path_mapping['soul']['path'],
            'memories': self.path_mapping['memories']['path'],
            'skills': self.path_mapping['skills']['path'],
            'sessions': self.path_mapping['sessions']['path'],
            'state/state.db': self.path_mapping['state_db']['path'],
            'logs': self.path_mapping['logs']['path'],
            'cron': self.path_mapping['cron']['path'],
            'webui/.token': self.path_mapping['webui_token']['path'],
            'webui/db': Path('/home/appuser/.hermes-web-ui'),
            'image_cache': self.path_mapping['image_cache']['path'],
            'baoyu_skills': self.path_mapping['baoyu_skills']['path'],
        }
        
        if not skip_restore:
            restore_mapping['config/config.yaml'] = self.path_mapping['config']['path']
        else:
            restored_path = self.hermes_home / 'config.yaml.restored'
            src = download_dir / 'config' / 'config.yaml'
            if src.exists():
                shutil.copy2(src, restored_path)
                logger.info("Restored config.yaml to config.yaml.restored for merge")
        
        for src_rel, dst in restore_mapping.items():
            if src_rel in skip_paths:
                logger.info(f"Skipping restore of {src_rel} (will be regenerated by entrypoint.sh)")
                continue
            
            src = download_dir / src_rel
            if src.exists():
                try:
                    if src.is_file():
                        dst.parent.mkdir(parents=True, exist_ok=True)
                        shutil.copy2(src, dst)
                        logger.info(f"Restored: {src_rel}")
                    elif src.is_dir():
                        if dst.exists():
                            shutil.rmtree(dst)
                        shutil.copytree(src, dst)
                        logger.info(f"Restored directory: {src_rel}")
                except Exception as e:
                    logger.error(f"Failed to restore {src_rel}: {e}")
            else:
                logger.warning(f"Not found in backup: {src_rel}")
        
        logger.success("Data restoration completed")


class ConfigFileHandler(FileSystemEventHandler):
    """Config file change handler  -  triggers immediate backup on change"""
    
    STARTUP_GRACE_PERIOD = 30
    
    def __init__(self, manager: DatasetManager):
        self.manager = manager
        self.last_backup_time = 0
        self.backup_cooldown = 10
        self.start_time = time.time()
        self._startup_logged = False
        
    def on_modified(self, event):
        """Called when a watched file is modified"""
        if event.is_directory:
            return
        
        elapsed = time.time() - self.start_time
        if elapsed < self.STARTUP_GRACE_PERIOD:
            if not self._startup_logged:
                logger.info(f"In startup grace period ({int(self.STARTUP_GRACE_PERIOD - elapsed)}s remaining), skipping backup for: {event.src_path}")
                self._startup_logged = True
            return
            
        watched_files = ['config.yaml', '.env', 'auth.json']
        if any(event.src_path.endswith(f) for f in watched_files):
            current_time = time.time()
            if current_time - self.last_backup_time > self.backup_cooldown:
                logger.info(f"Config file changed: {event.src_path}")
                logger.info("Triggering immediate backup...")
                try:
                    self.manager.upload_to_dataset()
                    self.last_backup_time = current_time
                    logger.success("Immediate backup completed")
                    self._trigger_reload()
                except Exception as e:
                    logger.error(f"Immediate backup failed: {e}")
    
    def _trigger_reload(self):
        """Attempt to trigger Hermes config reload"""
        logger.info("Configuration saved. Please restart Space to apply changes immediately.")


def run_daemon():
    """Background daemon mode  -  scheduled sync + real-time file watching"""
    logger.info("Starting data sync daemon...")
    
    default_interval = 300
    sync_interval = int(os.environ.get('SYNC_INTERVAL', str(default_interval)))
    
    if os.environ.get('SYNC_INCREMENTAL', 'true').lower() in ('true', '1', 'yes'):
        if sync_interval < 120:
            logger.warning(f"SYNC_INTERVAL={sync_interval}s is too short for incremental mode, consider >= 120s")
    
    manager = DatasetManager()
    
    if not manager.validate():
        logger.error("Configuration invalid, exiting")
        sys.exit(1)
    
    logger.info(f"Sync interval: {sync_interval} seconds")
    logger.info(f"Incremental backup: {manager.enable_incremental}")
    
    observer = None
    if WATCHDOG_AVAILABLE:
        try:
            logger.info("Starting file watcher for real-time sync...")
            event_handler = ConfigFileHandler(manager)
            observer = Observer()
            observer.schedule(event_handler, str(manager.hermes_home), recursive=False)
            observer.start()
            logger.success("File watcher started - config changes will trigger immediate backup")
        except Exception as e:
            logger.error(f"Failed to start file watcher: {e}")
            logger.warning("Falling back to scheduled sync only")
            observer = None
    else:
        logger.warning("Watchdog not available, using scheduled sync only")
    
    try:
        while True:
            try:
                time.sleep(sync_interval)
                logger.info("Performing scheduled backup...")
                manager.upload_to_dataset()
            except KeyboardInterrupt:
                logger.info("Daemon stopped")
                break
            except Exception as e:
                logger.error(f"Sync error: {e}")
    finally:
        if observer:
            logger.info("Stopping file watcher...")
            observer.stop()
            observer.join()
            logger.info("File watcher stopped")


def main():
    parser = argparse.ArgumentParser(description='Hermes Agent Data Sync')
    parser.add_argument('action', choices=['backup', 'restore', 'daemon'],
                       help='Action to perform')
    parser.add_argument('--force', '-f', action='store_true',
                       help='Force backup even if no changes')
    
    args = parser.parse_args()
    
    manager = DatasetManager()
    
    if not manager.validate():
        logger.error("Configuration invalid")
        sys.exit(1)
    
    if args.action == 'backup':
        success = manager.upload_to_dataset(force=args.force)
        sys.exit(0 if success else 1)
    elif args.action == 'restore':
        success = manager.download_from_dataset()
        sys.exit(0 if success else 1)
    elif args.action == 'daemon':
        run_daemon()


if __name__ == '__main__':
    main()