import gradio as gr
from typing import Dict, Any
from data_loader import read_dask_data, read_polars_data, read_another_dask_data

def load_and_process_data(dataset_choice: str, num_rows: int) -> Dict[str, Any]:
    """
    Load and process data based on the user's choice and specified number of rows.

    Args:
        dataset_choice (str): The dataset to load.
        num_rows (int): The number of rows to display.

    Returns:
        Dict[str, Any]: A dictionary containing the processed data or error message.
    """
    dataset_mapping = {
        "Dask Data": read_dask_data,
        "Polars Data": read_polars_data,
        "Another Dask Data": read_another_dask_data
    }

    data_loader = dataset_mapping.get(dataset_choice)
    if not data_loader:
        return {"error": "Invalid dataset choice."}

    try:
        data = data_loader()
        processed_data = data.head(num_rows)
        return {"processed_data": processed_data.to_dict()}
    except Exception as e:
        return {"error": f"Data processing failed: {str(e)}"}

def create_interface():
    """
    Create and launch the Gradio interface for data selection and display.
    """
    with gr.Blocks() as demo:
        gr.Markdown("# Enhanced Dataset Loader Demo")
        gr.Markdown("Interact with various datasets and select the amount of data to display.")

        with gr.Row():
            dataset_choice = gr.Dropdown(
                choices=["Dask Data", "Polars Data", "Another Dask Data"],
                label="Select Dataset"
            )
            num_rows = gr.Slider(
                minimum=1, maximum=100, value=5, label="Number of Rows to Display"
            )

        processed_data_output = gr.JSON(label="Processed Data")

        dataset_choice.render()
        num_rows.render()
        processed_data_output.render()

        demo.add(dataset_choice, num_rows, processed_data_output, load_and_process_data)

    demo.launch()

if __name__ == "__main__":
    create_interface()