import gradio as gr
import os
from git_commd import GitCommandWrapper
from typing import List, Optional

HF_TOKEN=os.environ['HF_TOKEN'] if 'HF_TOKEN' in os.environ else ''
WiseModel_TOKEN=os.environ['WM_TOKEN'] if 'WM_TOKEN' in os.environ else ''

def get_cache_dir():
    from random_word import RandomWords
    r = RandomWords()
    return r.get_random_word()


def check_disk():
    import os 
    return os.system("df -h /")  


def pull_from_wisemodel(token: str,
              url: str,
              repo_name: str,
               cache_dir,
              branch: Optional[str] = None):

    print("pull_from_wisemodel start")
    print(cache_dir)
    os.makedirs(cache_dir, exist_ok=True)
    # os.system("cd "+cache_dir)  
    gitCmd=GitCommandWrapper()
    isLfs=gitCmd.is_lfs_installed()
    if (not isLfs):
       gitCmd.git_lfs_install()
    gitCmd.clone(cache_dir,token,url,repo_name,branch)
    print("pull_from_wisemodel end")
    return f'Pulled {branch} to temp folder {cache_dir}: {url}'


def remove_file(cache_dir, repo_name):
    import os
    try:
        os.remove(f'{cache_dir}/{repo_name}')
    except:
        return ''
    return 'README.md file removed'
    

def push_to_hf(cache_dir, WiseModel_repo_name, hf_repo_id):
    from huggingface_hub import HfApi
    
    if not HF_TOKEN:
        raise gr.Error("Please enter your HF_TOKEN")
    print("push_to_hf start")
    api = HfApi(token=HF_TOKEN) # Token is not persisted on the machine.
    output = api.upload_folder(
        folder_path=f"{cache_dir}/{WiseModel_repo_name}",
        repo_id=hf_repo_id,
        repo_type="model",
    )
    print("push_to_hf end")
    return f'Pushed to {hf_repo_id}'
    
    
def handle(wisemodel_link, hf_repo_id):
    cache_dir = get_cache_dir()
    wiseModel_repo_url=wisemodel_link.replace(".git","").replace("git","").replace("clone","").replace(" ","")
    wiseModel_repo_info=wisemodel_link.replace(".git","").replace("git","").replace("clone","").replace(" ","").split("/")
    print(wiseModel_repo_info)
    wisemodel_repo_name=wiseModel_repo_info[len(wiseModel_repo_info)-1]
    stages = [
        (check_disk, (), {}),
        # # Run all the sanity checks on README.md
      
        (pull_from_wisemodel, (WiseModel_TOKEN,wiseModel_repo_url,wisemodel_repo_name, cache_dir,"main"), {}),
        (remove_file, (wisemodel_repo_name, cache_dir), {}),
        (check_disk, (), {}),
        (push_to_hf, (cache_dir, wisemodel_repo_name, hf_repo_id), {}),
        (check_disk, (), {}),
    ]

    results = []
    errors = []
    for func, args, kwargs in stages:
        try:
            results.append(str(func(*args, **kwargs)))
        except Exception as e:
            errors.append(str(e))
        if errors:
            break

    return '\n\n'.join(results), '\n\n'.join(errors)


with gr.Blocks() as demo:
    gr.Markdown('''


# wisemodel-to-HF/推送wisemodel上的模型、数据集等到Huggingface

# 这是一个示例Space,实际使用请参考下面说明，先Duplicate一个私有的space

- 这个space可以把已经发布在wisemodel的模型、数据集推送到Huggingface的相应repo。
- This space uploads model from WiseModel to Huggingface.

- **请确认您是huggingface上相应repo的创建者，拥有repo的写入的权限，否则会报错!**
- **Please make sure that you're the owner of the repo or have permission from the owner to do so!**


# 如何使用这个空间？
# How to use this Space?

- 点击右上角settings后面的“…”按钮，选择“Duplicate this Space”创建一个私有的space，同时输入wisemodel的token(公开模型选填，私有模型必填)，以及Huggingface的token（必填），确保有相应repo写入的权限。
- Duplicate this Space and providing WiseModel token (optional) and your read/write HF token (mandatory).

- 在HF上创建目标repo，这一步是必须的！
- Create your target model repo on HF. This step needs to be done manually. The Space doesn't do create an empty repo for you.

- 在刚刚自己创建的私有space里填写相应的信息，wisemodel的git clone链接，以及Huggingface的repo名称。
- In your own private Space, fill in information below.

- 点击submit按钮，然后可以通过logs按钮查看进度。
- Click submit then watch for output in container log for progress.

- Create README.md file (since the metadata is not compatible with HF)
    ''')
    
    wisemodel_link = gr.Textbox(label="Copy the git download link from the model detail page of wisemodel（从wisemodel上获取该模型的完整git clone链接） ")
    hf_repo_id = gr.Textbox(label="Target HF Model Repo ID (case sensitive). \nPlease make sure that this model has already been created")
    with gr.Row():
        button = gr.Button("Submit", variant="primary")
        clear = gr.Button("Clear")
    error = gr.Textbox(label="Error")
    output = gr.Textbox(label="Output")

    button.click(handle, [wisemodel_link, hf_repo_id], [output, error])

if __name__ == "__main__":
    demo.launch(debug = True)