import os
import sys
import tarfile
from huggingface_hub import HfApi, hf_hub_download

api = HfApi()
repo_id = os.getenv("HF_DATASET")
token = os.getenv("HF_TOKEN")
FILENAME = "latest_backup.tar.gz"

def restore():
    try:
        if not repo_id or not token:
            print("Skip Restore: HF_DATASET or HF_TOKEN not set")
            return
        print(f"Downloading {FILENAME} from {repo_id}...")
        path = hf_hub_download(
            repo_id=repo_id,
            filename=FILENAME,
            repo_type="dataset",
            token=token
        )
        with tarfile.open(path, "r:gz") as tar:
            tar.extractall(path="/root/.openclaw/")
        print(f"Success: Restored from {FILENAME}")
        return True
    except Exception as e:
        print(f"Restore Note: No existing backup found or error: {e}")

def backup():
    try:
        if not repo_id or not token:
            print("Skip Backup: HF_DATASET or HF_TOKEN not set")
            return
        # 要排除的路径（不备份的目录，使用相对路径）
        exclude = [
            "credentials",   # 避免泄露密钥
            "logs",          # 日志文件，无需备份
        ]
        with tarfile.open(FILENAME, "w:gz") as tar:
            # 遍历整个 /root/.openclaw，排除指定目录
            for root, dirs, files in os.walk("/root/.openclaw"):
                # 提前过滤排除目录
                dirs[:] = [d for d in dirs if d not in exclude]
                for file in files:
                    full_path = os.path.join(root, file)
                    # 计算 tar 包内的路径（去掉前缀 /root/.openclaw/）
                    arcname = full_path.replace("/root/.openclaw/", "")
                    tar.add(full_path, arcname=arcname)
        # 上传并覆盖
        api.upload_file(
            path_or_fileobj=FILENAME,
            path_in_repo=FILENAME,
            repo_id=repo_id,
            repo_type="dataset",
            token=token
        )
        print(f"Backup {FILENAME} Success (Overwritten).")
    except Exception as e:
        print(f"Backup Error: {e}")

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