src/data_sync.py

#!/usr/bin/env python3
"""
Hermes Agent Data Sync Service
Handles data persistence to/from Hugging Face Dataset strictly on SYNC_INTERVAL schedule
"""

import os
import sys
import time
import json
import shutil
import hashlib
import argparse
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


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.webui_home = Path(os.environ.get('HERMES_WEB_UI_HOME', '/data/.hermes-web-ui'))
        self.temp_dir = Path('/tmp/hermes_sync')
        self.last_sync_hash = None
        
        # Directories that represent REAL persistent user data
        self.backup_dirs = [
            'memories',
            'skills',
            'sessions',
            'cron',
            'script',
            'plugins',
            'extensions',
            'prompts',
            'knowledge',
        ]
        
        # Files that represent REAL persistent user config
        self.backup_files = [
            'config.yaml',
            '.env',
            'auth.json',
            'SOUL.md',
            'state.db',
            'channel_directory.json',
        ]
        
        # STRICT exclusion: never commit volatile logs, locks, or temporary files
        self.exclude_patterns = {
            '__pycache__', '.git', 'node_modules', '.cache', 'tmp', 'reports', 'src',
            'logs', 'gateway.lock', 'gateway.pid', 'gateway_state.json',
            '*.log', '*.db-shm', '*.db-wal', 'bridge.log', 'server.log',
            '.login-lock.json', 'preview-action.log'
        }
        
    def validate(self) -> bool:
        if not self.dataset_repo:
            logger.error("HF_DATASET_REPO not set")
            return False
        if not self.token:
            logger.warning("HF_TOKEN not set")
        return True

    def calculate_data_hash(self) -> str:
        """Calculate hash of actual user data (sessions, memories, skills, config, webui users)"""
        hasher = hashlib.md5()
        
        for fn in self.backup_files:
            fp = self.hermes_home / fn
            if fp.exists() and fp.is_file():
                try:
                    stat = fp.stat()
                    hasher.update(f"{fn}:{stat.st_size}:{stat.st_mtime}".encode())
                except:
                    pass
                    
        for dname in self.backup_dirs:
            dp = self.hermes_home / dname
            if dp.exists():
                for root, _, files in os.walk(dp):
                    for f in sorted(files):
                        if any(f.endswith(ext) for ext in ['.log', '.lock', '.pid', '.tmp', '.json']):
                            # session request dumps / temporary jsons shouldn't trigger commit loops
                            if 'request_dump' in f:
                                continue
                        try:
                            stat = os.stat(os.path.join(root, f))
                            hasher.update(f"{f}:{stat.st_size}:{stat.st_mtime}".encode())
                        except:
                            pass

        webui_db = self.webui_home / 'hermes-web-ui.db'
        if webui_db.exists():
            try:
                stat = webui_db.stat()
                hasher.update(f"webui_db:{stat.st_size}".encode())
            except:
                pass
                            
        return hasher.hexdigest()
    
    def prepare_backup_data(self) -> Path:
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
        self.temp_dir.mkdir(parents=True)
        
        for dirname in self.backup_dirs:
            (self.temp_dir / dirname).mkdir(exist_ok=True)
        
        def should_exclude(src: str, names: list) -> bool:
            ignored = []
            for n in names:
                if n in self.exclude_patterns or any(n.endswith(ext) for ext in ['.log', '.lock', '.pid', '.db-shm', '.db-wal']):
                    ignored.append(n)
                if 'request_dump' in n:
                    ignored.append(n)
            return ignored
        
        try:
            for dirname in self.backup_dirs:
                src_dir = self.hermes_home / dirname
                if src_dir.exists():
                    shutil.copytree(src_dir, self.temp_dir / dirname, dirs_exist_ok=True, ignore=should_exclude)
            
            for filename in self.backup_files:
                src_file = self.hermes_home / filename
                if src_file.exists():
                    shutil.copy2(src_file, self.temp_dir / filename)
            
            # Backup Web UI sqlite db and tokens only (exclude logs/locks)
            if self.webui_home.exists():
                webui_dst = self.temp_dir / 'webui_data'
                webui_dst.mkdir(parents=True, exist_ok=True)
                for f in os.listdir(self.webui_home):
                    if f.endswith(('.db', '.token')) and not f.endswith(('.db-shm', '.db-wal', '.lock.json')):
                        src_p = self.webui_home / f
                        if src_p.is_file():
                            shutil.copy2(src_p, webui_dst / f)

            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:
        try:
            current_hash = self.calculate_data_hash()
            if not force and self.last_sync_hash is not None and current_hash == self.last_sync_hash:
                logger.info("No data changes detected. Skipping backup upload.")
                return True

            backup_dir = self.prepare_backup_data()
            logger.info(f"Uploading scheduled backup to dataset: {self.dataset_repo}")
            
            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')}"
            )
            
            # Squash commit history so the dataset repo always stays at 1 commit
            try:
                self.api.super_squash_history(
                    repo_id=self.dataset_repo,
                    repo_type="dataset",
                    commit_message=f"Hermes Agent backup - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
                )
            except Exception as se:
                logger.warning(f"Squash history failed: {se}")
            
            self.last_sync_hash = current_hash
            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:
        try:
            logger.info(f"Downloading from dataset: {self.dataset_repo}")
            download_dir = Path('/tmp/hermes_download')
            if download_dir.exists():
                shutil.rmtree(download_dir)
            download_dir.mkdir(parents=True)
            
            self.api.snapshot_download(
                repo_id=self.dataset_repo,
                repo_type="dataset",
                local_dir=str(download_dir)
            )
            
            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):
        self.hermes_home.mkdir(parents=True, exist_ok=True)
        self.webui_home.mkdir(parents=True, exist_ok=True)
        
        skip_restore = os.environ.get('SKIP_CONFIG_RESTORE', 'true').lower() in ('true', '1', 'yes')
        
        restore_list = []
        for dirname in self.backup_dirs:
            restore_list.append((dirname, self.hermes_home / dirname))
        for filename in self.backup_files:
            restore_list.append((filename, self.hermes_home / filename))
        
        if not skip_restore:
            restore_list.append(('config.yaml', self.hermes_home / 'config.yaml'))
        else:
            restored_path = self.hermes_home / 'config.yaml.restored'
            src = download_dir / 'config.yaml'
            if src.exists():
                shutil.copy2(src, restored_path)
        
        for src_rel, dst in restore_list:
            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)
                    elif src.is_dir():
                        if dst.exists():
                            shutil.rmtree(dst)
                        shutil.copytree(src, dst)
                except Exception as e:
                    logger.error(f"Failed to restore {src_rel}: {e}")

        webui_src = download_dir / 'webui_data'
        if webui_src.exists():
            shutil.copytree(webui_src, self.webui_home, dirs_exist_ok=True)
        
        logger.success("Data restoration completed")


def run_daemon():
    logger.info("Starting data sync daemon...")
    sync_interval = int(os.environ.get('SYNC_INTERVAL', '1800'))
    manager = DatasetManager()
    
    if not manager.validate():
        logger.error("Configuration invalid, exiting")
        sys.exit(1)
    
    logger.info(f"Sync interval strictly set to: {sync_interval} seconds")
    
    try:
        while True:
            time.sleep(sync_interval)
            # Only commits on the strict schedule, AND only if real user data changed!
            manager.upload_to_dataset(force=False)
    except KeyboardInterrupt:
        pass


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

if __name__ == '__main__':
    main()