#!/usr/bin/env python3
import os, shutil, argparse
from huggingface_hub import Repository, create_repo

def deploy(hf_user, space_name, local_dir, private=False):
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise SystemExit("Please set HF_TOKEN env variable.")
    repo_id = f"{hf_user}/{space_name}"
    create_repo(repo_id=repo_id, repo_type="space", space_sdk="docker", private=private, exist_ok=True, token=token)
    clone_dir = space_name
    if os.path.exists(clone_dir):
        shutil.rmtree(clone_dir)
    repo = Repository(local_dir=clone_dir, clone_from=f"https://huggingface.co/spaces/{repo_id}", use_auth_token=token)
    for item in os.listdir(local_dir):
        s = os.path.join(local_dir, item)
        d = os.path.join(clone_dir, item)
        if os.path.isdir(s):
            if os.path.exists(d): shutil.rmtree(d)
            shutil.copytree(s, d)
        else:
            shutil.copy2(s, d)
    with open(os.path.join(clone_dir, ".gitattributes"), "a") as f:
        f.write("\n*.pkl filter=lfs diff=lfs merge=lfs -text\n")
    repo.git_add()
    try: repo.git_commit("Deploy Space")
    except: pass
    repo.push_to_hub()
    print("Deployed -> https://huggingface.co/spaces/"+repo_id)

if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--hf_user', required=True)
    p.add_argument('--space_name', required=True)
    p.add_argument('--local_dir', required=True)
    p.add_argument('--private', action='store_true')
    args = p.parse_args()
    deploy(args.hf_user, args.space_name, args.local_dir, private=args.private)