# Import the necessary libraries:
# - `gradio` is a library for creating interactive web interfaces
# - `typing` provides type annotations for Python
# - `data_loader` is a custom module that contains functions to read different types of data

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

# Define a function called `load_and_process_data` that takes a dataset choice and the number of rows to display
# This function is responsible for loading and processing the data based on the user's input
def load_and_process_data(dataset_choice: str, num_rows: int) -> Dict[str, Any]:
    try:
        # Create a mapping of dataset choices to their corresponding data loading functions
        dataset_mapping = {
            "Dask Data": read_dask_data,
            "Polars Data": read_polars_data,
            "Another Dask Data": read_another_dask_data
        }

        # Fetch the appropriate data loading function based on the user's dataset choice
        data_loader = dataset_mapping.get(dataset_choice)
        if not data_loader:
            # If the dataset choice is invalid, return an error message
            return {"error": "Invalid dataset choice."}

        # Load the data using the selected data loading function
        data = data_loader()

        # Process the data to show the specified number of rows
        processed_data = data.head(num_rows)

        # Convert the processed data to a dictionary for JSON serialization
        return {
            "processed_data": processed_data.to_dict()
        }
    except Exception as e:
        # If an exception occurs during data processing, log the error
        # and return an error message
        print(f"Error processing data: {str(e)}")
        return {"error": "Unable to process data. Please check the logs for details."}

# Define a function called `create_interface` that creates a Gradio interface
# The interface allows the user to select a dataset and the number of rows to display
def create_interface():
    # Define the input components for the interface
    dataset_choice = gr.components.Dropdown(
        choices=["Dask Data", "Polars Data", "Another Dask Data"], 
        label="Select Dataset"
    )
    num_rows = gr.components.Slider(
        minimum=1, maximum=100, value=5, label="Number of Rows to Display"
    )

    # Define the layout of the Gradio interface
    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.render()
            num_rows.render()
        processed_data_output = gr.JSON(label="Processed Data")
        processed_data_output.render()

    # Add the input components and the data processing function to the interface
    demo.add(dataset_choice, num_rows, processed_data_output, load_and_process_data)

    # Launch the Gradio interface
    demo.launch()

# Execute the `create_interface` function when the script is run
if __name__ == "__main__":
    create_interface()