#!/usr/bin/env python3
import os
import sys

# Redirect HuggingFace cache to /tmp to avoid permission issues inside /config
os.environ["HF_HOME"] = "/tmp/hf_home"
os.environ["HF_HUB_CACHE"] = "/tmp/hf_home/hub"
os.environ["XDG_CACHE_HOME"] = "/tmp"

from huggingface_hub import HfApi, create_repo, upload_folder, snapshot_download

TOKEN = os.environ.get("HF_TOKEN")
api = HfApi(token=TOKEN)
try:
    username = api.whoami()['name']
except Exception:
    username = "Lorilbjr"

REPO_ID = f"{username}/antigravity-data"
SYNC_DIR = "/config"

import subprocess

def save():
    print(f"Compressing {SYNC_DIR} to /tmp/backup.tar.gz (this may take a moment)...")
    try:
        # Create dataset repo if it doesn't exist
        create_repo(repo_id=REPO_ID, repo_type="dataset", token=TOKEN, exist_ok=True)
        
        # Compress /config, excluding .cache
        cmd = ["tar", "-czf", "/tmp/backup.tar.gz", "--exclude=.cache", "-C", SYNC_DIR, "."]
        subprocess.run(cmd, check=True)
        
        print(f"Uploading backup.tar.gz to Hugging Face...")
        from huggingface_hub import upload_file
        upload_file(
            path_or_fileobj="/tmp/backup.tar.gz",
            path_in_repo="backup.tar.gz",
            repo_id=REPO_ID,
            repo_type="dataset",
            token=TOKEN
        )
        print("Successfully saved to HF!")
        
        # Cleanup
        if os.path.exists("/tmp/backup.tar.gz"):
            os.remove("/tmp/backup.tar.gz")
            
    except subprocess.CalledProcessError as e:
        print(f"Error compressing files: {e}")
    except Exception as e:
        print(f"Error saving: {e}")

def load():
    import shutil
    print(f"Downloading backup.tar.gz from Hugging Face...")
    try:
        from huggingface_hub import hf_hub_download
        
        if not os.path.exists(SYNC_DIR):
            os.makedirs(SYNC_DIR)
            
        file_path = hf_hub_download(
            repo_id=REPO_ID,
            filename="backup.tar.gz",
            repo_type="dataset",
            token=TOKEN,
            cache_dir="/tmp/hf_home/hub",
            local_dir="/tmp"
        )
        
        print(f"Extracting to {SYNC_DIR}...")
        cmd = ["tar", "-xzf", file_path, "-C", SYNC_DIR]
        subprocess.run(cmd, check=True)
        
        print("Successfully loaded from HF!")
        
        # Cleanup
        if os.path.exists(file_path):
            os.remove(file_path)
            
    except subprocess.CalledProcessError as e:
        print(f"Error extracting files: {e}")
    except Exception as e:
        print(f"Error loading: {e} (Maybe no backup exists yet?)")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 sync.py [save|load]")
    elif sys.argv[1] == "save":
        save()
    elif sys.argv[1] == "load":
        load()