import pandas as pd
from openpyxl import load_workbook
import gradio as gr
import os

# Function to extract and map data from the input workbook
def transform_data(input_path, mapping_df):
    # Load the input workbook
    input_workbook = pd.ExcelFile(input_path)
    
    # Initialize a dictionary to store data for output
    output_data = {col: [] for col in mapping_df["PO Output Column"] if not pd.isna(col)}

    # Iterate through each mapping rule
    for _, row in mapping_df.iterrows():
        output_column = row["PO Output Column"]
        input_rule = row["UVM  MMB  POLY STICKER Column"]

        if pd.isna(output_column) or pd.isna(input_rule):
            continue

        # Handle fixed values
        if "Fixed" in input_rule:
            fixed_value = input_rule.split("-")[0].strip()
            output_data[output_column] = [fixed_value] * 1  # Placeholder for rows

        # TODO: Add logic to map specific columns from input workbook

    return pd.DataFrame(output_data)

# Main processing function
def process_files(input_workbook, mapping_file, output_template):
    try:
        # Load the data mapping file
        mapping_df = pd.read_excel(mapping_file, sheet_name="Output")

        # Transform the data
        transformed_data = transform_data(input_workbook, mapping_df)

        # Load the output template
        output_workbook = load_workbook(output_template)
        output_sheet = output_workbook["363040"]

        # Write transformed data to the output sheet
        for row_idx, row_data in enumerate(transformed_data.itertuples(index=False), start=2):
            for col_idx, value in enumerate(row_data, start=1):
                output_sheet.cell(row=row_idx, column=col_idx, value=value)

        # Save the generated output file
        output_file_path = "Generated_Output.xlsx"
        output_workbook.save(output_file_path)

        return output_file_path

    except Exception as e:
        return f"An error occurred: {e}"

# Define the Gradio interface
def generate_excel(input_workbook, mapping_file, output_template):
    result = process_files(input_workbook.name, mapping_file.name, output_template.name)
    if os.path.exists(result):
        return result
    else:
        return "An error occurred during file generation."

with gr.Blocks() as app:
    gr.Markdown("# Excel Sheet Generator")
    gr.Markdown("Upload the input workbook, mapping file, and output template to generate the final Excel file.")

    with gr.Row():
        input_workbook = gr.File(label="Input Workbook", file_types=[".xlsx"])
        mapping_file = gr.File(label="Mapping File", file_types=[".xlsx"])
        output_template = gr.File(label="Output Template", file_types=[".xlsx"])

    output_file = gr.File(label="Generated Excel File")
    generate_button = gr.Button("Generate Excel File")

    generate_button.click(generate_excel, inputs=[input_workbook, mapping_file, output_template], outputs=[output_file])

# Launch the app
app.launch()