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"

BASE_PATH = "/root/.openclaw/"

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

        print("Restore success")
    except Exception as e:
        print(f"No backup found or restore failed: {e}")


def backup():
    try:
        if not repo_id or not token:
            print("Skip Backup: HF_DATASET or HF_TOKEN not set")
            return

        with tarfile.open(FILENAME, "w:gz") as tar:
            paths = [
                "/root/.openclaw/sessions",
                "/root/.openclaw/agents/main/sessions",
                "/root/.openclaw/openclaw.json"
            ]

            for p in paths:
                if os.path.exists(p):
                    arcname = p.replace(BASE_PATH, "")
                    tar.add(p, arcname=arcname)

        api.upload_file(
            path_or_fileobj=FILENAME,
            path_in_repo=FILENAME,
            repo_id=repo_id,
            repo_type="dataset",
            token=token
        )

        print("Backup success")
    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()