import sys
import time
import os
from huggingface_hub import snapshot_download, HfApi

REPO_ID = "Airpyk98/hermes-ui-storage"
LOCAL_DIR = os.path.expanduser("~/.hermes")

def pull():
    try:
        print("Pulling dataset from HF...", flush=True)
        os.makedirs(LOCAL_DIR, exist_ok=True)
        snapshot_download(repo_id=REPO_ID, repo_type="dataset", local_dir=LOCAL_DIR)
        print("Dataset pulled successfully.", flush=True)
    except Exception as e:
        print(f"Error pulling dataset: {e}", flush=True)

def push():
    try:
        if not os.path.exists(LOCAL_DIR) or not os.listdir(LOCAL_DIR):
            return
        api = HfApi()
        api.upload_folder(
            folder_path=LOCAL_DIR,
            repo_id=REPO_ID,
            repo_type="dataset",
            commit_message="Auto-sync from Workspace",
            commit_description="Syncing memory and config."
        )
        print("Dataset pushed successfully.", flush=True)
    except Exception as e:
        print(f"Error pushing dataset: {e}", flush=True)

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "pull":
        pull()
    elif len(sys.argv) > 1 and sys.argv[1] == "push":
        push()
    elif len(sys.argv) > 1 and sys.argv[1] == "daemon":
        while True:
            time.sleep(300) # Every 5 minutes
            push()