import os
import shutil
import subprocess
import signal
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
import gradio as gr

from huggingface_hub import HfApi
from huggingface_hub import ModelCard

from gradio_huggingfacehub_search import HuggingfaceHubSearch

from apscheduler.schedulers.background import BackgroundScheduler

HF_TOKEN = os.environ.get("HF_TOKEN")

HF_PATH = "https://huggingface.co/"

def button_click(hf_model_id, conv_template, quantization):
    api = HfApi()
    model_dir_name = hf_model_id.split("/")[1]
    mlc_model_name = model_dir_name + "-" + quantization + "-" + "MLC"

    os.system("mkdir -p dist/models")
    os.system("git lfs install")

    api.snapshot_download(repo_id=hf_model_id, local_dir=f"./dist/models/{model_dir_name}")

    os.system("mlc_llm convert_weight ./dist/models/" + model_dir_name + "/" + \
              " --quantization " + quantization + \
              " -o dist/" + mlc_model_name)
    
    os.system("mlc_llm gen_config ./dist/models/" + model_dir_name + "/" + \
              " --quantization " + quantization + " --conv-template " + conv_template + \
              " -o dist/" + mlc_model_name + "/")
    
    # push to HF
    user_name = api.whoami()["name"]
    api.create_repo(repo_id=f"{user_name}/{mlc_model_name}", private=True)

    api.upload_large_folder(folder_path=f"./dist/{mlc_model_name}",
                  repo_id=f"{user_name}/{mlc_model_name}",
                  repo_type="model")
    
    return "successful"

demo = gr.Interface(
    fn=button_click,
    inputs = [gr.Textbox(label="HF Model ID"),
              gr.Dropdown(["tinyllama_v1_0", "qwen2"], label="Conversation Template"),
              gr.Dropdown(["q4f16_1", "q4f32_1"], label="Quantization Method")],
    outputs = "text"
)

demo.launch()