##################################### 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_file):
    # Here you can process the uploaded file (dataset_file) as needed
    return f"Hello {name}!! Using model: {model_name} with template: {prompt_template}"

##################################### App UI #######################################

# Function to build the interface based on user choice
def build_interface(choice):
    if choice == "Predefined Dataset":
        dataset_input = gr.Dropdown(label="Predefined Dataset", choices=['1', '2', '3'], value='1', key="dataset_predefined")
    elif choice == "Upload Your Own":
        dataset_input = gr.File(label="Upload Dataset", accept=".csv,.txt", key="dataset_upload")
    else:
        dataset_input = None
    return dataset_input

# Function to handle changes in dataset choice
def on_choice_change(choice):
    interface = build_interface(choice)
    update_interface(interface)

# Function to update the interface with new dataset input
def update_interface(dataset_input):
    demo.Interface(
        fn=greet,
        inputs=[model_name, prompt_template, name_input, dataset_input],
        outputs=output
    ).launch()

##################################### Gradio Blocks #######################################
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")

    # Initial interface setup
    initial_choice = dataset_choice.value
    initial_interface = build_interface(initial_choice)

    # Output textbox
    output = gr.Textbox(label="Output")

    # Setup button
    tune_btn = gr.Button("Start Fine Tuning")
    tune_btn.click(on_choice_change, inputs=[dataset_choice])

    # Launch the initial interface
    update_interface(initial_interface)

    # Update visibility based on user choice
    dataset_choice.change(on_choice_change, inputs=[dataset_choice])

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

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