##################################### Imports ######################################
# Generic imports
import gradio as gr
import json
import os

########################### Global objects and functions ###########################

def get_json_cfg():
    """Retrieve configuration file"""
    config_path = os.getenv('CONFIG_PATH')
    with open(config_path, 'r') as file:
        config = json.load(file)
    return config

conf = get_json_cfg()

def greet(model_name, prompt_template, name, dataset):
    return f"Hello {name}!! Using model: {model_name} with template: {prompt_template}"

##################################### App UI #######################################
with gr.Blocks() as demo:

    ##### Title Block #####
    gr.Markdown("# Instruction Tuning with Unsloth")

    ##### Model Inputs #####

    # Select Model
    model_name = gr.Dropdown(label="Model", choices=conf['model']['choices'], value="gpt2")
    # Prompt template
    prompt_template = gr.Textbox(label="Prompt Template", value="Instruction: {0}\nOutput: {1}")
    # Prompt Input
    name_input = gr.Textbox(label="Your Name")
    # Dataset choice
    dataset_choice = gr.Radio(label="Choose Dataset", choices=["Predefined Dataset", "Upload Your Own"], value="Predefined Dataset")
    dataset_predefined = gr.Dropdown(label="Predefined Dataset", choices=['1', '2', '3'], value='1', visible=True)
    dataset_upload = gr.UploadButton(label="Upload Dataset", file_types=[".pdf",".csv",".jsonl"], visible=False) # gr.File(label="Upload Dataset", visible=False)

    

    


    
    def dropdown_visibility(radio):  # Accept the event argument, even if not used
        value = radio  # Get the selected value from the radio button
        if value == "Predefined Dataset":
            return gr.Dropdown(visible=bool(1))
        else:
            return gr.Dropdown(visible=bool(0))

    def upload_visibility(radio):
        value = radio
        if value == "Upload Your Own":
            return gr.UploadButton(visible=bool(1)) #make it visible
        else:
            return gr.UploadButton(visible=bool(0))

    dataset_choice.change(dropdown_visibility, dataset_choice, dataset_predefined)
    dataset_choice.change(upload_visibility, dataset_choice, dataset_upload)

    # Define function to update visibility based on dataset choice
    def update_visibility():
        choice = dataset_choice.value
        if choice == "Predefined Dataset":
            dataset_predefined.visible = True
            dataset_upload.visible = False
        elif choice == "Upload Your Own":
            dataset_predefined.visible = False
            dataset_upload.visible = True

    # Bind visibility update function to change event of dataset_choice
    dataset_choice.change(update_visibility)
    
    # Update visibility based on user choice
    #dataset_predefined, dataset_upload = dataset_choice.change(update_dataset_visibility, inputs=[dataset_choice], outputs=[dataset_predefined, dataset_upload])
    
    ##### Model Outputs #####

    # Text output
    output = gr.Textbox(label="Output")
    
    ##### Execution #####

    # Setup button
    tune_btn = gr.Button("Start Fine Tuning")
    # Execute button
    tune_btn.click(fn=greet, 
                   inputs=[model_name, prompt_template, name_input, dataset_predefined],
                   outputs=output)

##################################### Launch #######################################

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