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

# 文件监控（可选）
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'
        
        # 增量备份配置（可通过环境变量控制）
        self.enable_incremental = os.environ.get('SYNC_INCREMENTAL', 'true').lower() in ('true', '1', 'yes')
        self.sync_interval_override = os.environ.get('SYNC_INTERVAL', None)
        
        # 数据路径映射
        # skip_backup=True 表示该路径在增量模式下跳过备份（大文件/临时文件）
        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},  # 日志大且不重要
            'cron': {'path': self.hermes_home / 'cron', 'skip_backup': False},
            'webui_token': {'path': Path('/data/.hermes-web-ui') / '.token', 'skip_backup': False},
            'image_cache': {'path': self.hermes_home / 'image_cache', 'skip_backup': True},  # 图片缓存大
            'baoyu_skills': {'path': Path('/home/appuser/.baoyu-skills'), 'skip_backup': False},
        }
        
    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, will try public dataset")
            
        return True
    
    def _calculate_file_hash(self, file_path: Path) -> str:
        """计算文件 SHA256 哈希（只读前 128KB 加速）"""
        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:
        """计算目录中所有文件的组合哈希"""
        if not dir_path.exists():
            return ""
        try:
            sha256 = hashlib.sha256()
            for root, dirs, files in os.walk(dir_path):
                # 跳过隐藏文件
                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:
        """加载上次同步状态"""
        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):
        """保存同步状态"""
        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:
        """检查是否有文件变化（增量备份核心）"""
        if not state.get('last_sync'):
            return True  # 首次同步
        
        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:
        """准备备份数据到临时目录（支持增量备份）"""
        logger.info("Preparing backup data...")
        
        # 增量备份检查 - 如果没有变化则跳过
        if self.enable_incremental:
            state = self._load_sync_state()
            if not self._has_changes(state):
                logger.info("No changes detected, skipping backup (incremental mode)")
                # 更新最后同步时间
                state['last_sync'] = time.time()
                self._save_sync_state(state)
                return Path('/dev/null')  # 特殊标记，upload_to_dataset 会检测到
        
        # 清理并创建临时目录
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
        self.temp_dir.mkdir(parents=True)
        
        # 创建目录结构（不再创建 logs/ 和 image_cache/）
        (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()
        
        # 复制文件
        try:
            # 配置文件
            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')
            
            # 环境变量（敏感信息）- 上传前 masking 处理
            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 Hugging Face tokens to prevent upload rejection
                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)
                with open(env_dst, 'w', encoding='utf-8') as f:
                    f.write(env_content)
            
            # OAuth 认证
            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')
            
            # 人格定义
            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')
            
            # 记忆
            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)
            
            # 技能
            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)
            
            # 会话
            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)
            
            # 数据库
            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')
            
            # 日志 - 跳过（大且不重要）
            # if self.path_mapping['logs'].exists():
            #     shutil.copytree(self.path_mapping['logs'], self.temp_dir / 'logs', dirs_exist_ok=True)
            
            # 定时任务
            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 - 跳过（大文件）
            # 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 用户配置
            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 认证 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')
            
            # 添加元数据
            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:
        """上传数据到 Hugging Face Dataset"""
        try:
            backup_dir = self.prepare_backup_data()
            
            # 检查是否是增量备份的跳过标记
            if backup_dir == Path('/dev/null'):
                logger.info("Backup skipped (no changes)")
                return True
            
            logger.info(f"Uploading to dataset: {self.dataset_repo}")
            
            # 上传文件夹到 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')}"
            )
            
            # 更新同步状态
            if self.enable_incremental:
                state = self._load_sync_state()
                state['last_sync'] = time.time()
                # 更新文件哈希
                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:
        """从 Hugging Face Dataset 下载数据"""
        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)
            )
            
            logger.success("Download completed")
            
            # 恢复数据到 Hermes 目录
            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):
        """从下载的目录恢复数据
        
        注意: config.yaml 在恢复时被跳过，因为 entrypoint.sh 会根据环境变量
        重新生成正确的 config.yaml。如果恢复旧的 config.yaml，会导致模型
        配置被覆盖（例如 minimaxai/minimax-m2.7 被替换为旧模型）。
        """
        logger.info("Restoring data to Hermes home...")
        
        # 确保目标目录存在
        self.hermes_home.mkdir(parents=True, exist_ok=True)
        
        # 恢复策略控制
        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'],
            '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):
    """配置文件变化处理器 - 实时同步到 Dataset 并触发重载"""
    
    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):
        """文件被修改时触发"""
        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):
        """尝试触发 Hermes 配置重载"""
        logger.info("Configuration saved. Please restart Space to apply changes immediately.")


def run_daemon():
    """后台守护进程模式 - 定期同步 + 实时文件监听"""
    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()