import gradio as gr
import pandas as pd
from data_labelling import main as label_data

def preview_data(file):
    if file is None:
        return "Please upload a CSV file."
    
    df = pd.read_csv(file.name)
    preview = df.head()  # Show the first few lines of the dataset
    return gr.DataFrame(preview), df.shape  # Return the preview and the shape of the dataset

def run_labeling(input_file, output_file):
    # Run the labeling script
    labeled_file = label_data(input_file, output_file)
    return labeled_file  # Return the path to the labeled file

def main_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# YouTube Comments Labeling")
        
        with gr.Row():
            file_input = gr.File(label="Upload your CSV file")
            preview_button = gr.Button("Preview Data")
        
        data_preview = gr.DataFrame()
        data_shape = gr.Textbox(label="Dataset Shape", interactive=False)
        
        preview_button.click(preview_data, inputs=file_input, outputs=[data_preview, data_shape])
        
        with gr.Row():
            label_option = gr.Radio(["No", "Yes"], label="Do you want to label your dataset?", value="No")
        
        output_name = gr.Textbox(label="Output File Name (with .csv extension)", value="labeled_dataset.csv")
        label_button = gr.Button("Run Labeling")
        
        labeled_file_path = gr.Textbox(label="Labeled File Path", interactive=False)
        download_button = gr.File(label="Download Labeled File")
        
        def handle_labeling(input_file, label_choice, output_name):
            if label_choice == "Yes":
                labeled_file = run_labeling(input_file.name, output_name)
                return labeled_file, labeled_file
            else:
                return None, None
        
        label_button.click(
            handle_labeling, 
            inputs=[file_input, label_option, output_name], 
            outputs=[labeled_file_path, download_button]
        )
    
    demo.launch()

if __name__ == "__main__":
    main_interface()