import gradio as gr
from huggingface_hub import HfApi
import os

hf_api = HfApi()
CURRENT_MODEL_PATH = "./current_model"

HF_TOKEN = os.environ.get("HF_TOKEN", None)
REPO = os.getenv("REPO")
REVISION = os.getenv("REVISION")
TARGET = os.getenv("TARGET")

def func(hf_token, repo, revision, target_repo):
    hf_token = hf_token if hf_token else HF_TOKEN
    repo = repo if repo else REPO
    revision = revision if revision else REVISION
    target_repo = target_repo if target_repo else TARGET
    
    print('download the desired model locally ...')
    hf_api.snapshot_download(
        repo_id=repo,
        revision=revision,
        token=hf_token,
        local_dir=CURRENT_MODEL_PATH,
    )
    print(os.listdir(CURRENT_MODEL_PATH))

    print(f"The content of {repo}/{revision} has been copied to {target_repo}!" )

    for file in os.listdir(CURRENT_MODEL_PATH):
        print(f'uploading {file}')
        with open(f"{CURRENT_MODEL_PATH}/{file}", "rb") as fobj:
            hf_api.upload_file(
                path_or_fileobj=fobj,
                path_in_repo=f"{file}",
                repo_id=target_repo,
                repo_type="model",
                token=hf_token,
                commit_message=f"updating model files: {file}-{revision}",
            )
    
    return f"The content of {repo}/{revision} has been copied to {target_repo}!" 

iface = gr.Interface(fn=func, inputs=["text","text","text"], outputs="text")
iface.launch()