import pandas as pd
import gradio as gr
from ui_components import CustomDropdown, CustomButton  # Example: Replace with actual component names

# Function to load Excel sheets and extract first column data from valid sheets
def load_excel_sheets(file_path):
    # Load all sheets into a dictionary {sheet_name: DataFrame}
    xls = pd.ExcelFile(file_path)
    sheets = {}
    
    for sheet_name in xls.sheet_names:
        df = pd.read_excel(xls, sheet_name=sheet_name)
        # Check if the sheet has at least one column and is not empty
        if not df.empty and df.shape[1] > 0:
            # Extract the first column and keep the order intact
            sheets[sheet_name] = df.iloc[:, 0].dropna().tolist()
    
    return sheets

# Function to get the order of sheets from the main sheet
def get_sheet_order(file_path, main_sheet_name):
    # Read the main sheet which contains the sheet order
    df_order = pd.read_excel(file_path, sheet_name=main_sheet_name)
    # Assuming the sheet names are listed in the first column
    return df_order.iloc[:, 0].dropna().tolist()

# Function to create the combined prompt from selected values and weights
def combine_selected_items(*selected_items_and_weights):
    combined_prompt = []
    for i in range(0, len(selected_items_and_weights), 2):
        item = selected_items_and_weights[i]
        weight = selected_items_and_weights[i+1]
        if item:
            combined_prompt.append(f"{item} (Weight: {weight})" if weight else item)
    return f"Combined Prompt: {' | '.join(combined_prompt)}"

# Function to create the Gradio interface
def prompt_generator_interface(file_path, main_sheet_name):
    # Load all the sheets and extract data from the first column
    sheets = load_excel_sheets(file_path)

    # Get the sheet order from the main sheet
    sheet_order = get_sheet_order(file_path, main_sheet_name)

    # Filter out sheets that don't exist in the main order
    valid_sheets = {sheet_name: sheets[sheet_name] for sheet_name in sheet_order if sheet_name in sheets}

    # Define the sheets that should NOT have weights
    no_weight_sheets = ["Poetic", "Scenarios", "Camera Setup", "Combos", "PositiceNegative prompts", "LAZY Mode"]

    # Move Resources tab to the last place
    if 'Resources' in valid_sheets:
        valid_sheets['Resources'] = valid_sheets.pop('Resources')

    # Gradio interface using custom UI components
    with gr.Blocks() as interface:
        gr.Markdown("# 📝 Witness Prompt Generator\nSelect an item from each sheet to generate a combined prompt.")

        # Initialize an empty list to store all the dropdowns and weight inputs
        dropdowns_and_weights = []

        # Display all sheets as dropdowns in the order specified by the main sheet
        with gr.Row():
            for sheet_name in valid_sheets:
                with gr.Column():
                    # Using the custom dropdown from ui_components.py
                    dropdown = CustomDropdown(choices=valid_sheets[sheet_name], label=sheet_name, interactive=True)
                    dropdowns_and_weights.append(dropdown)
                    
                    # If the sheet is not in the no_weight_sheets, add a weight input
                    if sheet_name not in no_weight_sheets:
                        # Custom weight input
                        weight_input = gr.Textbox(label=f"{sheet_name} Weight", placeholder="Enter weight (optional)", interactive=True)
                        dropdowns_and_weights.append(weight_input)

        # Using a custom button for submission
        submit_button = CustomButton("Generate Combined Prompt")

        # Textbox to display the combined prompt
        combined_output = gr.Textbox(label="Combined Prompt", placeholder="Your combined prompt will appear here...")

        # Action when the submit button is clicked
        submit_button.click(combine_selected_items, inputs=dropdowns_and_weights, outputs=combined_output)

    return interface

# Create and launch the Gradio interface
if __name__ == "__main__":
    # Path to your Excel file (adjust the path if necessary)
    file_path = 'Witness Prompt Generator.xlsm'
    
    # Define the name of the main sheet that contains the order of sheets
    main_sheet_name = 'Main Order'  # Change this to match the name of your control sheet
    
    # Create and launch interface
    interface = prompt_generator_interface(file_path, main_sheet_name)
    interface.launch()