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"
OPENCLAW_DIR = "/root/.openclaw"

# 备份时跳过的文件/目录模式
EXCLUDE_SUFFIXES = (".log", ".tmp", ".bak", ".pid")
EXCLUDE_NAMES = {"__pycache__", "node_modules", ".git"}


def should_exclude(tarinfo):
    """过滤掉不需要备份的临时文件"""
    name = os.path.basename(tarinfo.name)
    if name in EXCLUDE_NAMES:
        return None
    if any(name.endswith(s) for s in EXCLUDE_SUFFIXES):
        return None
    return tarinfo


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=OPENCLAW_DIR)

        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

        if not os.path.exists(OPENCLAW_DIR):
            print(f"Skip Backup: {OPENCLAW_DIR} does not exist")
            return

        with tarfile.open(FILENAME, "w:gz") as tar:
            tar.add(OPENCLAW_DIR, arcname=".", filter=should_exclude)

        size_mb = os.path.getsize(FILENAME) / (1024 * 1024)
        print(f"Backup archive size: {size_mb:.1f} MB")

        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()