import gradio as gr
import os
import zipfile
import glob
from huggingface_hub import login, HfApi, create_repo

def export_model_to_hf(hftoken, experiment_name, manual_epoch_number, logs_path, repoid, create_new_repo):
    num_epochs = int(manual_epoch_number) if manual_epoch_number.isdigit() else None
    
    # Construct the weights path based on the provided number of epochs
    if num_epochs is not None:
        weights_path = f"/content/RVC/assets/weights/{experiment_name}_e{num_epochs}*"
    else:
        potential = f"/content/RVC/assets/weights/{experiment_name}.pth"
        if os.path.exists(potential):
            weights_path = f"/content/RVC/assets/weights/{experiment_name}"
        else:
            currentMax = 0
            for r, _, f in os.walk("/content/RVC/assets/weights/"):
                for name in f:
                    if(name.endswith(".pth") and (name != experiment_name + ".pth")):
                        if(name.find(experiment_name) == -1):
                            continue
                        pot = name.split('_')
                        ep = pot[len(pot) - 2][1:]
                        if not ep.isdecimal():
                            continue
                        ep = int(ep)
                        if ep > currentMax:
                            currentMax = ep
            num_epochs = currentMax
            weights_path = f"/content/RVC/assets/weights/{experiment_name}_e{num_epochs}*"
    
    weights_files = glob.glob(weights_path + ".pth")
    if weights_files and any(glob_result := glob.glob(logs_path)):
        log_file = glob_result[0]
        output_folder = "/content/toHF"
        os.makedirs(output_folder, exist_ok=True)
        output_zip_path = f"{output_folder}/{experiment_name}.zip"
        with zipfile.ZipFile(output_zip_path, 'w') as zipf:
            for weights_file in weights_files:
                zipf.write(weights_file, os.path.basename(weights_file))
            zipf.write(log_file, os.path.basename(log_file))
        
        login(token=hftoken)
        if create_new_repo:
            create_repo(repoid)
        
        api = HfApi()
        api.upload_folder(folder_path=output_folder, repo_id=repoid, repo_type="model")
        
        return f"Model uploaded successfully to {repoid}"
    else:
        return "Couldn't find your model files. Check the found file results above. (Did you run Index Training?)"

with gr.Blocks() as demo:
    gr.Markdown("# Export Finished Model to HuggingFace test<br>[click this to get HF token](https://huggingface.co/settings/tokens)</small>")
    
    hftoken = gr.Textbox(label="HuggingFace Token (set Role to 'write')", type="password")
    experiment_name = gr.Textbox(label="Experiment Name", value="rewrite")
    manual_epoch_number = gr.Textbox(label="Manual Epoch Number (leave blank for auto-detect)", value="")
    logs_path = gr.Textbox(label="Logs Path", value="/content/RVC/logs/rewrite/added_IVF37_Flat_nprobe_1_rewrite_v2.index")
    repoid = gr.Textbox(label="HuggingFace Repository ID", value="Hev832/rewrite-sonic")
    create_new_repo = gr.Checkbox(label="Create New Repository", value=True)
    
    output = gr.Textbox(label="Output")
    
    btn = gr.Button("Export Model")
    btn.click(export_model_to_hf, inputs=[hftoken, experiment_name, manual_epoch_number, logs_path, repoid, create_new_repo], outputs=output)

demo.launch()