import gradio as gr
import os
from datasets import load_dataset, Dataset
import pandas as pd
from huggingface_hub import login

def load_huggingface_dataset(dataset_link, token):
    """
    Load a Hugging Face dataset using the provided link and token.
    """
    # Extract dataset name and config if applicable
    parts = dataset_link.split("/")
    dataset_name = parts[-1]
    if len(parts) > 2:
        owner = parts[-2]
        dataset_name = f"{owner}/{dataset_name}"
    
    # Load the dataset
    dataset = load_dataset(dataset_link, split="train")
    
    # Return the dataset as a DataFrame with index and transcription columns
    df = dataset.to_pandas().reset_index()
    return df[["index", "transcription"]], dataset

def update_transcriptions(df, dataset, token,dataset_link):
    """
    Update the transcriptions in the dataset and push it back to the Hugging Face Hub.
    """
    # Convert DataFrame back to Dataset
    updated_dataset = Dataset.from_pandas(df)
    
    # print(updated_dataset)
    # print(dataset)
    
    # Replace the original transcription column in the dataset
    dataset = dataset.map(
        lambda examples, idx: {"transcription": updated_dataset["transcription"][idx]},
        with_indices=True 
    )
    print(dataset['transcription'][0])
    
    login(token)
    
    dataset.push_to_hub(dataset_link)
    
    return "Dataset updated and changes submitted to the Hugging Face Hub!"

# Gradio Interface
def main():
    
    
    dataset = None  # To store the loaded dataset object globally
    
    def load_dataset_and_show_table(dataset_link, token):
        """
        Load the dataset and return the DataFrame to display in Gradio.
        """
        nonlocal dataset
        df, dataset = load_huggingface_dataset(dataset_link, token)
        return df
    
    def submit_changes(df, token,dataset_link):
        """
        Submit updated changes to the Hugging Face Hub.
        """
        if dataset is None:
            return "No dataset loaded to update."
        
        print(df)
        print(token)
        return update_transcriptions(df, dataset, token,dataset_link)

    # Gradio Interface
    with gr.Blocks(css=".dataframe-row { height: 200px; }") as interface:
        gr.Markdown("## Hugging Face Audio Dataset Editor")
        
        # Input fields for dataset link and token
        dataset_link = gr.Textbox(label="Hugging Face Dataset Link")
        hf_token = gr.Textbox(label="Hugging Face Token", type="password")
        
        # Button to load dataset
        load_button = gr.Button("Load Dataset")
        
        # Table to display and edit dataset
        table = gr.Dataframe(
            headers=["index", "transcription"],
            datatype=["number", "str"],
            interactive=True,
            label="Edit Dataset (Transcriptions are RTL)",
        )
        
        # Button to submit changes
        submit_button = gr.Button("Submit Changes")
        output_message = gr.Textbox(label="Message")
        
        # RTL styling for transcription column
        table.style = {"transcription": {"direction": "rtl"}}
        
        # Button functionality
        load_button.click(load_dataset_and_show_table, [dataset_link, hf_token], table)
        submit_button.click(submit_changes, [table, hf_token,dataset_link], output_message)
    
    # Launch Gradio Interface
    interface.launch(share=True)
    
    

if __name__ == "__main__":
    main()