#!/usr/bin/env python3
import os
import sys
import tarfile
from huggingface_hub import HfApi, hf_hub_download

# 配置常量
BACKUP_DIR = "/root/.openclaw"
BACKUP_FILENAME = "backup.tar.gz"

def restore():
    repo_id = os.getenv("HF_DATASET")
    token = os.getenv("HF_TOKEN")
    
    if not repo_id or not token:
        print("--- [SYNC] 跳过恢复: 未配置 HF_DATASET 或 HF_TOKEN ---")
        return True # 返回 True 表示流程继续，只是没恢复数据

    try:
        print(f"--- [SYNC] 启动恢复流程, 目标仓库: {repo_id} ---")
        api = HfApi()
        
        # 获取文件列表
        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset", token=token)
        
        if BACKUP_FILENAME  in files:
            print(f"--- [SYNC] 发现备份文件: {BACKUP_FILENAME}, 正在下载... ---")
            # 下载到本地临时位置
            downloaded_path = hf_hub_download(
                repo_id=repo_id, 
                filename=BACKUP_FILENAME, 
                repo_type="dataset", 
                token=token,
                force_download=True # 确保每次都是最新的
            )
            
            extract_path = BACKUP_DIR
            print(f"--- [SYNC] 正在解压至 {extract_path} ... ---")
            
            with tarfile.open(downloaded_path, "r:gz") as tar:
                # 安全解压逻辑 (兼容旧版本 Python)
                def is_within_directory(directory, target):
                    abs_directory = os.path.abspath(directory)
                    abs_target = os.path.abspath(target)
                    prefix = os.path.commonprefix([abs_directory, abs_target])
                    return prefix == abs_directory

                for member in tar.getmembers():
                    member_path = os.path.join(extract_path, member.name)
                    if not is_within_directory(extract_path, member_path):
                        raise Exception(f"尝试跳出目录的攻击行为: {member.name}")
                
                tar.extractall(path=extract_path)
            
            print(f"--- [SYNC] 恢复成功! 数据已覆盖至 {extract_path} ---")
            return True
        
        else:
            # 修正日志：不再谎称检查了时间
            print("--- [SYNC] 未在仓库中找到备份文件，将使用空状态启动 ---")
            return True
    except Exception as e:
        print(f"--- [SYNC] 恢复异常: {e} ---")
        import traceback
        traceback.print_exc()
        return True # 即使失败也返回 True，避免阻塞启动

def backup():
    repo_id = os.getenv("HF_DATASET")
    token = os.getenv("HF_TOKEN")

    if not repo_id or not token:
        print("--- [SYNC] 跳过备份: 未配置 HF_DATASET 或 HF_TOKEN ---")
        return

    try:
        print(f"--- [SYNC] 正在执行全量备份: {BACKUP_FILENAME} ---")
        
        base_path = BACKUP_DIR
        if not os.path.exists(base_path):
            print(f"--- [SYNC] 警告: 备份目录 {base_path} 不存在，跳过打包 ---")
            # 即使目录不存在，也可以上传一个空的标记文件或者直接返回，这里选择返回
            return
        
        # 创建临时压缩包
        with tarfile.open(BACKUP_FILENAME, "w:gz") as tar:
            # 修复点：将 arcname 设置为 '.'，这样解压后直接是 .openclaw 的内容，
            # 或者设置为 os.path.basename(base_path) 保留文件夹名。
            # 通常为了还原方便，我们希望在 .openclaw 目录下直接看到内容，所以用 '.'
            # 如果希望解压出来是 /root/.openclaw/.openclaw/... 则用 basename
            arcname = "." 
            tar.add(base_path, arcname=arcname)
            print(f"--- [SYNC] 本地打包完成，大小: {os.path.getsize(BACKUP_FILENAME) / 1024 / 1024:.2f} MB ---")
        
        api = HfApi()
        print(f"--- [SYNC] 正在上传至 Hugging Face ({repo_id}) ... ---")
        
        api.upload_file(
            path_or_fileobj=BACKUP_FILENAME, 
            path_in_repo=BACKUP_FILENAME, 
            repo_id=repo_id, 
            repo_type="dataset", 
            token=token
        )
        
        print(f"--- [SYNC] 备份上传成功! ---")
        
        # 清理本地临时备份文件以节省空间
        if os.path.exists(BACKUP_FILENAME):
            os.remove(BACKUP_FILENAME)
            print("--- [SYNC] 本地临时文件已清理 ---")
            
    except Exception as e:
        print(f"--- [SYNC] 备份失败: {e} ---")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "backup":
        backup()
    else:
        restore()